torus711 のアレ

主に競技プログラミングの問題について書きます.PC 以外だと数式が表示されないかもしれないです

AtCoder Beginner Contest 359, D : Avoid K Palindrome

問題概要

 A, B, ? からなる長さ $n$ の文字列 $S$ が与えられる.$S$ 中の ? を(それぞれ独立に)A または B に置き換えてできる文字列であって,長さ $k$ の任意の(連続する)部分文字列が回文でないようなものの個数を ${} \bmod {998{,}244{,}353}$ で求めよ.

制約

  • $2 \leq k \leq n \leq 1{,}000$
  • $k \leq 10$

解法

 $\bmod$ をとることについては,適切に処理するなり AtCoder Library (ACL) を使うなりすればよいので割愛します.
 空文字列から始めて,末尾に文字を追加することで合法な文字列を構成していくことを考えます.今手元に $i$ 文字の文字列があるとして末尾に連結できる文字は,回文に関することを一旦置いておけば,

  • $S_{ i + 1 } \neq $ B のとき,A を追加できる.
  • $S_{ i + 1 } \neq $ A のとき,B を追加できる.

となります.回文に関する条件を考慮するには,

  • 連結後の文字列が $k$ 文字未満である.または,
  • 末尾 $k$ 文字が回文でない.

を満たさないものを弾くようにすればよいです.
 これで問題が解けました……とはならなくて,最悪ケースのとき(具体的には全部 ? な $S$)に作れる文字列の個数が $\Theta( 2^n )$ 個になってやばいので,一工夫必要です.
 上で挙げた条件をよく見ると,末尾付近の高々 $k$ 文字にしか依存していません.これは,末尾付近の $k$ 文字が同じ文字列に対しては,以降で可能な連結の仕方が同じであるということです.このことを利用して,
\begin{equation*}
\mathit{ dp }[i][s] := \text{$i$ 文字からなる文字列で,末尾 $k$ 文字が $s$ で表されるものの個数}
\end{equation*}
と状態をとる DP を考えることができます.ここで,末尾 $k$ 文字が $s$ で表されるというのは,$s$ が(少なくとも)$k$ ビットからなる整数で,($2$ 進で)$s$ の右から $j$ 桁目を $s_j$ と書くことにすると,
\begin{align*}
s_j = \begin{cases}
1 & \text{(構成中の文字列が $k$ 文字以上で,右から $i$ 文字目が B)} \\
0 & \text{(otherwise)}
\end{cases}
\end{align*}
であることを言います*1*2. この DP の初期状態は
\begin{equation*}
\mathit{ dp }[i][s] = \begin{cases}
1 & \text{($i = 0, s = 0$)} \\
0 & \text{(otherwise)}
\end{cases}
\end{equation*}
となります.$\mathit{ dp }[i][s]$ からの遷移は,$c \in \{$ A$, $ B$\}$ に対して,
\begin{align*}
dp[ i + 1 ][ s' ] \overset{+}{\leftarrow} dp[i][s] &&
\text{if $i + 1 < k \lor \lnot p( k, s )$}
\end{align*}
と更新します*3.ここで,$s'$ は $s, k, c$ から計算できる値で,

  1. $s$ を $1$ ビット左シフトし,
  2. 右 $k$ ビットを取り出し,
  3. $c = $ B なら $1$,そうでないなら $0$ との bitwise OR をとる.

とすることで求まる値です.$p( k, s )$ は $s$ の下位 $k$ ビットが回文になっているかどうかを判定する関数です.中身は下記の実装例をご参照ください.
 この DP の状態数は $\Theta( n 2^k )$ 個です.各状態からの遷移は $O( 1 )$ 個ですが,遷移可能かどうかを判定するために $O( k )$ 時間がかかります.なので全体としては $O( n k 2^k )$ 時間となります.

コード

#include <atcoder/modint>
using namespace atcoder;
using MINT = modint998244353;

bool palindrome( const int K, const int s )
{
	REP( i, K )
	{
		if ( !!( s & 1 << i ) != !!( s & 1 << ( K - 1 - i ) ) )
		{
			return false;
		}
	}
	return true;
}

int main()
{
	IN( int, N, K );
	IN( string, S );

	static MINT dp[ 1024 ][ 1024 ];
	// dp[ length of string ][ histtory of the last K ] := ways
	dp[0][0] = 1;

	REP( i, N )
	{
		REP( j, 1 << K )
		{
			FOR( c, "AB"s )
			{
				if ( c == 'A' && S[i] != 'B' ||
						c == 'B' && S[i] != 'A' )
				{
					const int nj = ( j << 1 ) & ( ( 1 << K ) - 1 ) | ( c == 'B' );
					if ( i + 1 < K || !palindrome( K, nj ) )
					{
						dp[ i + 1 ][ nj ] += dp[i][j];
					}
				}
			}
		}
	}

	cout << accumulate( ALL( dp[N] ), MINT() ).val() << endl;

	return 0;
}

*1:右からとか左からとか,$1$ に対応付けるのが A とか B とかのバリエーションはあります.

*2:equation 環境内で code タグを使いたい人生だった…….

*3:$\overset{+}{\leftarrow}$ は左辺に右辺を加算代入することを表すとします.