2020年12月26日土曜日

CodeChef - Counting Strings

問題: 与えられた文字列$s$について、左から比較しても右から比較しても辞書順で大きい文字列$t$を数える。

$t$の構成を考えると、2つの位置$i, j (i \le j)$を境界にして、以下のように左から順に5つの部分に分解できる。(2番目と4番目は1文字)

  • $s$と完全一致
  • $s[i] < t[i]$
  • 文字自由
  • $s[j] < t[j]$
  • $s$と完全一致

$N := |s|$とし、0-basedで考える。$f(i) :=$ 文字$s[i]$より大きい文字の数($s[i] = \operatorname{a}$なら$f(i)=25$)とすると、$i, j$が異なる場合の数え上げは$\sum_{i < j} f(i)f(j)26^{j-i-1}$である。あとは$i, j$が一致する場合、つまり$s$と$t$がちょうど1字異なる場合を別に数えて$f(0)+...+f(N-1)$を足す。

上の式を見ると、畳み込める形になっている。$g(x) = f(N-x-1)$として一方の添え字を逆順にすると、$26^{N-2}\sum_{i+j \le N-2} f(i)g(j) (1/26)^{i+j}$と変形できて、$f$と$g$をFFTで畳み込んだ後で$1/26$を掛けながら加算すればよい。$\bmod$は$10^9+7$だが、登場する数が十分に小さい(高々$25^2 \cdot 10^5$)ので$998244353$でやっても問題ない。


よく見ると、変形して累積和を取るだけだった。

$$ \sum_{i < j} f(i)f(j)26^{j-i-1} = \frac{1}{26}\sum_{i=0}^{N-1} \biggl( \frac {f(i)}{26^i} \sum_{j = i+1}^{N-1}{f(j)}26^j \biggr ) $$

形を見て畳み込みできそうと思えるのは悪くないとしても、変数を分離できないかまず考えるほうがさすがによい気がする。