SRM428 Div.1 500 The Long Palindrome - Darseinのプロコン日記

個人的に好きなタイプの問題だったので久しぶりに解説を書きます。

問題

アルファベット小文字('a'-'z')からなる任意の文字列を考える。
長さN以下、含まれる文字の種類数K以下の回文の個数を求めなさい。
ただし、数が大きくなりうるので1234567891の剰余を返しなさい。

制約: 1<=N<=10^9, 1<=K<=26

解法

まずは高速なアルゴリズムを一旦無視して、答えを数式にしてみる。
$p_{i,j}$ を、「長さがちょうどiで文字種類数がちょうどjの回文の個数」とする。
すると答えは、 $\displaystyle \sum_{i=1}^{N} \sum_{j=1}^{K}p_{i,j}$ となるので、これを求めたい。

ここで、回文は前半が決まってしまえば後半は一意に定まるという特徴を考えると、
n/2文字の任意の文字列を作れば、その対称をとることで回文は漏れなく重複無く数え上げることができる。
よって、 $s_{i,j}$ を、「長さがちょうどiで文字種類数がちょうどjの文字列の個数」とすると、
$\displaystyle \sum_{i=1}^{N} \sum_{j=1}^{K}p_{i,j} = \sum_{i=1}^{N} \sum_{j=1}^{K}s_{\lceil i/2 \rceil,j} = \sum_{i=1}^{\lceil N/2 \rceil} \sum_{j=1}^{K}s_{i,j} + \sum_{i=1}^{\lfloor N/2 \rfloor} \sum_{j=1}^{K}s_{i,j}$ と答えを書き直せる。

これで各 $s_{i,j}$ を求めればよい問題になったが、それぞれO(1)で求められたとしても全部でO(NK)かかるため間に合わない。
しかしそれはそれとして、O(NK)で求めるDPを考える。
これは、 $s_{i,j} = j \cdot s_{i-1,j} + (26-(j-1)) \cdot s_{i-1,j-1}$ という漸化式に基づくDPとなる。
式の意味は、前の項が「もし今まで使ったことがあるj文字のどれかをつかえば、新たに使った文字種類数が増えることはない」という意味で、後ろの項は「今まで使っていない(26-j)文字のうちどれかを使うと、文字種類数は1だけ増える」という意味である。
これを小さいi,jから順に計算して埋めていくことで、O(NK)で計算ができる。

しかし、前述の通り、これでは遅すぎなので、計算量を改善する。
ここで注目するのは、上記の漸化式にはiに依存する邪魔なパラメータは無く、すべてjのみに依存して更新されているという点だ。
この形のDPは、K×K行列を用いた行列積で、1つ次の項を求めることができる。これを利用すると、

によって、 $s_{N,K}$ をO(K^3 logN)で求められるようになる。

しかし、求めたいのは $s_{N,K}$ ではなく、その和である。
というようなことが蟻本の3-4にもそのまま書いてあるので、そのテクニックを用いる。
行列の累乗和を求めるのは、行列を4つ並べて2K×2Kの行列を作り、左上にA、右上に零行列、残り下半分をそれぞれ単位行列にして累乗計算することで求められる。詳しくは蟻本を読もう(布教活動)。
あとは、求めた累乗Sに(1, 0, ..., 0)^Tを掛けて出てくるベクトル要素の和が答え。modを取るのは忘れずに。