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 eitherAorBsuch 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: