NAC 2020 I

#include <bits/stdc++.h>
using namespace std;

const int MX = 60000, MOD = 998244353;

int m, mx, ans = 0, dp[2][MX];
string s;

vector<int> convert(int u) {
    vector<int> ve;
    for (int i = 0; i < s.size(); i++, u /= 3) {
        ve.push_back(u % 3 - 1);
    }
    reverse(ve.begin(), ve.end());
    return ve;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> s >> m;
    mx = pow(3, s.size());
    dp[0][mx - 1] = 1;
    for (int i = 1; i <= 20; i++) {
        int cur = i & 1, prv = cur ^ 1;
        fill(dp[cur], dp[cur] + MX, 0);
        for (int j = 0; j < mx; j++) {
            if (dp[prv][j] == 0) {
                continue;
            }
            vector<int> ve = convert(j);
            for (char c = 'A'; c <= 'Z'; c++) {
                int msk = 0, di = i - 1, up = i - 1, le = i;
                for (int k = 0; k < s.size(); k++) {
                    up += ve[k];
                    int at = min(min(up, le) + 1, di + (c != s[k]));
                    (msk *= 3) += at - le + 1;
                    di = up; le = at;
                }
                (dp[cur][msk] += dp[prv][j]) %= MOD;
                if (le == m) {
                    (ans += dp[prv][j]) %= MOD;
                }
            }
        }
    }
    cout << (ans + ((int)s.size() == m)) % MOD;
}