Official

G - Avoid K Palindrome Editorial by en_translator


For a string \(T\) consisting of A, B, and ?, define \(\operatorname{dp} _ T[P]\) as follows (DP stands for Dynamic Programming):

  • The number of strings obtained by replacing each occurrnce of ? with either A or B such that:
    • it is a good string, and
    • its last \(\min\lbrace K-1,|T|\rbrace\) characters coincide with \(P\).

For a string \(T\) consisting of A, B, and ?, suppose that \(\operatorname{dp} _ T[P]\) is known for all possible \(P\).
Then, if \(\operatorname{dp} _ {T+\mathtt{A}}[P]\) and \(\operatorname{dp} _ {T+\mathtt{B}}[P],\operatorname{dp} _ {T+\mathtt{?}}[P]\) can be found fast, then the problem can be solved using Dynamic Programming. (Here, the string \(X\) succeeded by a character \(c\) is denoted by \(X+c\).

Let us consider an algorithm to find \(\operatorname{dp} _ {T+\mathtt A}\) from \(\operatorname{dp} _ T\). (Same applies to B, and for ? we can take their sum.)

  • Initially, initialize \(\operatorname{dp} _ {T+\mathtt A}[P]=0\) for all \(P\).
  • For all possible \(P\), do the following.
    • If \(P+\mathtt A\) is a length-\(W\) palindrome, do nothing.
    • Otherwise, let \(P ^ \prime\) be the string formed by the last \(\min\lbrace K-1,|P|+1\rbrace\) characters of \(P+\mathtt A\). Add \(\operatorname{dp} _ T[P]\) to \(\operatorname{dp} _ {T+\mathtt A}[P ^ \prime]\).

Consider finding \(\operatorname{dp} _ S\) for the given string \(S\) using these values. Even if you spend \(\Theta(K)\) time to check if a string is palindrome, the overall complexity is \(O(2 ^ KKN)\), which is fast enough.

For implementation, one can use a string or a non-negative integer as the key of the DP table.

If you us a string, the key is easy to understand, and you can also use a sentinel to evaluate the DP table for string of length less than \(K\), but the constant factor may be bad.
If you use a non-negative integer, the constant factors of the time and spacial complexity is good, but the implementation may be complicated.

The following is sample code.

In the sample code in C++, the key of the DP table is a non-negative integer, and it has a separate process for shorter \(S\). In the sample code in Python, the key of the DP table is a string, and it puts a sentinel to simplify the implementation.

#include <algorithm>
#include <atcoder/modint>
#include <iostream>
#include <numeric>
#include <ranges>
#include <string>
#include <vector>

int main() {
    using namespace std;
    using modint = atcoder::static_modint<998244353>;

    unsigned N, K;
    cin >> N >> K;

    // Non-palindrome 0/1 string of length K
    // represented as a integer
    vector<unsigned> not_palindrome;
    for (unsigned i = 0; i < 1U << K; ++i) {
        bool is_palindrome = true;
        for (unsigned j = 0, k = K - 1; j < k; ++j, --k){
            // If there are different bits seen from the head or tail,
            if ((1 & i >> j) != (1 & i >> k)){
                // Then it's not a palindrome.
                is_palindrome = false;
                break;
            }
        }
        // Add if it is not a palindrome
        if (!is_palindrome)
            not_palindrome.emplace_back(i);
    }

    string S;
    cin >> S;

    vector<modint> dp(1U << K - 1), prev(1U << K - 1);
    { // Enumerate all possible first K-1 characters
        // a_mask = bitmask with bit 1 for the position 'A', and 0 otherwise
        // q_mask = bitmask with bit 0 for the position '?', and 1 otherwise
        unsigned a_mask{}, q_mask{};
        for (const auto c : S | views::take(K - 1)) {
            (a_mask *= 2) += c == 'A';
            (q_mask *= 2) += c != '?';
        }

        // Fix the positions where q_mask is standing, and enumerate the other part
        for (unsigned i{q_mask}; i < 1U << K - 1; ++i |= q_mask)
            dp[i ^ a_mask] = 1;
    }

    const unsigned mask{(1U << K - 1) - 1};
    for (const auto c : S | views::drop(K - 1)) {
        swap(dp, prev);
        ranges::fill(dp, modint{});

        // If 'A' is to be added
        if (c != 'B')
            // Transition occurs only when it is not a palindrome and it ends with 'A' (0)
            for(const auto i : not_palindrome | views::filter([](auto i){return ~i & 1;}))
                dp[i & mask] += prev[i / 2];

        // If 'B' is to be added
        if (c != 'A')
            // Transition occurs only when it is not a palindrome and it ends with 'B' (1)
            for(const auto i : not_palindrome | views::filter([](auto i){return i & 1;}))
                dp[i & mask] += prev[i / 2];
    }

    // The answer is the sum over all prefixes
    cout << reduce(begin(dp), end(dp)).val() << endl;
    return 0;
}
N, K = map(int, input().split())

S = input()

# Initially, fill with characters that are not A nor B
mp = {'C' * (K - 1) : 1}

for c in S:
    # Add one character to the tail
    # Merge the cases for adding 'A' and adding 'B'
    # If both addition is impossible, merge an empty dictionary
    tmp = ({s + 'A' : v for s, v in mp.items()} if c != 'B' else {}) | ({s + 'B' : v for s, v in mp.items()} if c != 'A' else {})
    
    mp = {} # Erase the DP table
    
    for s, v in tmp.items():
        if s != s[::-1]: # If it is not a palindrome, drop the first character and add it
            if s[1:] in mp:
                mp[s[1:]] += v
            else:
                mp[s[1:]] = v

# The sum modulo 998244353 is the answer
print(sum(mp.values()) % 998244353)

posted:
last update: