ABC358のEをPythonで解く方法（復習）

E-Alphabet Tilesの問題の解法が思いつかず&実装方法も思いつかず悔しかったのでこの記事で復習します。

問題文

AtCoder Landでは、アルファベットの書かれたタイルが販売されています。高橋君は、タイルを一列に並べてネームプレートを作ろうと考えました。

長さ$1$以上$K$以下の英大文字からなる文字列であって、以下の条件を満たすものの個数を998244353で割った余りを求めてください。

• $1 \leq i \leq 26$ を満たす任意の整数$i$について以下が成立する。
  - 辞書順で$i$番目の英大文字を$a_i$とおく。例えば、$a_1 = A, a_5 = E, a_{26} = Z$である。
  - 文字列の中に含まれている$a_i$の個数は0以上$C_i$個以下である。

失敗したコード

from collections import deque, defaultdict
import numpy as np

mod = 998244353

K = int(input())
C = list(map(int,input().split()))

check = defaultdict(lambda: 0)
keys = deque([[0]*len(C)])
check[tuple(keys[0])] = 1
#print(check)

for i in range(K):
    len_keys = len(keys)
    for j in range(len_keys):
        key = keys.popleft() # [0 or 1 for in range(26)]
        #print("key:",key)
        #print("tuple_key:",tuple(key))
        check_num = check[tuple(key)]
        #print("check_num:",check_num)
        can_add = np.where(np.array(key) < np.array(C))[0]
        #print("can_add:",can_add)
        for c in can_add:
            key_copy = key.copy()
            key_copy[c] += 1
            if check[tuple(key_copy)] == 0:
                keys.append(key_copy)
            check[tuple(key_copy)] += check_num
    #print(check)
    print(f"iter {i} finished,",len(keys))

s = 0
for key in check.keys():
    if sum(list(key)) > 0:
        s += check[key]
        s %= mod
print(s)

$i$番目のアルファベットが何個ある配列が何個あるかを数える(defaultdictで)っていうアイデアにまでは至ったが、動的計画法の実装には至らなかった。(探索数は削ることができるがオーダーはそんなに変わってないかも)

動的計画法だとx軸｛-,A,B,...,Z｝とy軸｛0,1,...,K｝で何通りかを数えなければならないが、そのような方法が思いつかなかった。

解法

$dp[i][j]$を{A,..,i番目までのアルファベット}でj文字の配列の数と捉える.

例えば入力例1のC=[2,1,1,0,...,0]でK=2である場合を考える。

1. 0文字の文字列は""のみなので、$dp[0][0] = dp[1][0] = ... = dp[26][0] = 1$
2. 1文字の文字列は$C_i$が0より大きい$i$に対応するアルファベット1文字のみの文字列となるので、この場合はA,B,Cの3つだけ。
3. 2文字の場合、$i-1$番目のアルファベットのみでできている0文字および1文字の文字列に新しく$i$番目のアルファベットを加えることを考える。ただし、このとき加える新しいアルファベットの文字数は$C_i$を超えない。例えばAとBで2文字の文字列を構築することを考えると、AA,AB,BAの3通りである。
   • Aで構築される0文字の文字列にBを2つ加える：$dp[1][0]\times\binom{2}{2}=1$通り
   • Aで構築される1文字の文字列にBを1つ加える： $dp[1][1]\times\binom{2}{1}=2$通り
   • Aのみで構築される2文字の文字列の数も包含：$dp[0][0]\times\binom{2}{2}=1$通り
   • したがって合計4通りとなる。
4. このようにi-1文字目までで構築されるk-c文字の文字列のにi文字目をc文字加える操作を繰り返せば求める文字列の数を求めることができる。

\begin{align}
dp[0][0] &= dp[1][0] = ... = dp[26][0] = 1\\
dp[0][1] &= ... = dp[0][K] = 0\\
dp[1][1] &= dp[0][1] + dp[0][0] \times \binom{1}{1} = 1\\
dp[2][1] &= dp[1][1] + dp[1][0] \times \binom{1}{1} = 2\\
dp[3][1] &= dp[2][1] + dp[2][0] \times \binom{1}{1} = 3\\
...\\
dp[1][2] &= dp[0][2] + dp[0][0] \times \binom{2}{2} = 1\\
dp[2][2] &= dp[1][2] + dp[1][0] \times \binom{2}{2} + dp[1][1] \times \binom{2}{1} = 4\\
dp[3][2] &= dp[2][2] + dp[2][0] \times \binom{2}{2} + dp[2][1] \times \binom{2}{1} = 7\\
...\\
answer &= \sum_{k=1}^{2} dp[3][k]
\end{align}

この処理を一般化すると、求める文字列の数は以下のように求めることができる。

\begin{align}
    dp[0][0] &= dp[1][0] = ... = dp[26][0] = 1\\
    dp[0][1] &= ... = dp[0][K] = 0\\
    dp[i>0][j>0] &= \left(\sum^{min(C_i,j)}_{c=0} dp[i-1][j-c] \times \binom{j}{c}\right) (mod \, 998244353)\\
    answer &= \sum^K_{k=1}dp[26][k]
\end{align}

実装する時の注意点

基本的な理論は以上になるが、これをそのまま実装しようとするとTLEになる。

from math import factorial
from functools import cache

mod = 998244353

# niko, niko_recursiveどっちでやってもTLEだった。
@cache
def niko(n,k): # math.factorialで階乗を計算させる
    return factorial(n) // (factorial(n-k)*factorial(k)) 

@cache
def niko_recursive(n,k): # パスカルの三角形を使う
    if n==k or k==0:
        return 1
    if n-k < k:
        return niko_recursive(n, n-k)
    else:
        return (niko_recursive(n-1,k-1)+niko_recursive(n-1,k)) % mod

K = int(input())
C = list(map(int,input().split()))

dp = [[0]*(K+1) for _ in range(len(C)+1)]
for i in range(len(C)+1):
    dp[i][0] = 1
    
s = 0
for k in range(1,K+1):
    for i in range(1,len(C)+1):
        for c in range(min(k,C[i-1])+1):
            #print(chr(ord('A')-1+i),c,k,dp[i-1][k-c] * niko(k,c))
            dp[i][k] = (dp[i][k] + (dp[i-1][k-c] * niko_recursive(k,c)) % mod) % mod
    s = (s + dp[len(C)][k]) % mod
print(s)

恐らく割り算、掛け算、再帰に時間がかかりすぎているため(オーダーがKlogK倍されてしまうため?)だと考えられる。これを解消するために必要なのが逆元という考え方だ。

逆元とは？

$a \times b \equiv 1 \pmod M $であるとき$b$は$a$の逆元となる。
つまり、$M$で割った余りを考えるときに$a$の逆数と同じような役割をしているのが逆元である。

この逆元を計算する方法としては主に2つの方法がある。

1. フェルマーの小定理を使う

フェルマーの小定理とは、$p$を素数、$a$を$p$と互いに素である整数だとすると以下の式が成り立つことである。

a^p \equiv a \pmod p

この式全体を$a^2$で割ると上式の右辺が逆元となる。
これはpythonのpow関数を用いて簡単に実装できるが、pythonのバージョンによっては$mod-2$乗の値を計算しなくてはならないように実装されているため、これでもなお計算に時間がかかることがある。

pow(a,-1,mod) # これで動けば一番楽
pow(a,mod-2,mod) # バージョン等によってはこの形式でしか計算できない、時間かかる

2. ユーグリッド互除法を応用する

ユーグリッド互除法とは、高校数学などで最小公倍数を求めるときや整数問題の応用で使った方法である。
例えば互いに素である$x,y$の最小公倍数を計算するとする。

\begin{align}
    x &= a_0 \times y + b_0 \\
    y &= a_1 \times b_0 + 1 \\
    y + b_0 \times (-a_1) &= 1 \\
    y + (x - (y \times a_0)) \times (-a_1)&= x \times (-a_1) + y \times (1 + a_0 \times a_1) \\
    &= x \times A + y \times B = 1
\end{align}

このように、ユーグリッド互除法を用いて$x \times A + y \times B = 1$という形式にすることができる。

このとき$x=998244353$であれば$y \times B \equiv 1 \pmod x$となるので$B$は$y$の逆元である。

def euclid(x,y):
    # x = y * a + b
    if x%y == 1:
        # x + y * (-a) = 1
        return 1, -(x//y)
    else:
        a_0, b_0 = euclid(y,x%y)
        # y * a_0 + (x - y * (x // y)) * b_0 = 1
        return b_0, a_0 - (x//y) * b_0

_, inv = euclid(mod, i)

この逆元の求め方を使えばACになった。

from functools import cache

mod = 998244353

K = int(input())
C = list(map(int,input().split()))

dp = [[0]*(K+1) for _ in range(len(C)+1)]
for i in range(len(C)+1):
    dp[i][0] = 1

#fact = [1] * (K+1)
#invfact = [1] * (K+1)
#for i in range(1,K+1):
#    fact[i] = fact[i-1] * i % mod
#    invfact[i] = pow(fact[i],mod-2,mod)

@cache
def euclid(x,y): # ユーグリッド互除法を用いて逆元を求める
    # x = y * a + b
    a = x // y
    b = x % y
    if b==1:
        # x + y * (-a) = 1
        return 1,-a
    else:
        a_, b_ = euclid(y,b)
        #y * a_ + b * b_ = 1
        #b = (x - y * a)
        #y * a_ + (x - y * a) * b_ = x * b_ + y * (a_ - a*b_) = 1
        return b_, a_ - a * b_
    
@cache
def fact(i):
    if i==0 or i==1:
        return 1
    return fact(i-1) * i % mod

@cache
def invfact(i):
    #return pow(fact(i),mod-2,mod)
    if i==1:
        return 1
    _, inv = euclid(mod,i)
    inv %= mod
    return invfact(i-1) * inv % mod

def comb(n,k):
    if n==k or k==0:
        return 1
    return fact(n) * invfact(n-k) % mod * invfact(k) % mod

s = 0
for k in range(1,K+1):
    for i in range(1,len(C)+1):
        for c in range(min(k,C[i-1])+1):
            #print(chr(ord('A')-1+i),c,k,dp[i-1][k-c] * niko(k,c))
            dp[i][k] = (dp[i][k] + (dp[i-1][k-c] * comb(k,c)) % mod) % mod
    s = (s + dp[len(C)][k]) % mod
print(s)

組み合わせの場合の数計算の方法について、組み合わせnCrの総和　計算の高速化及び逆元について(Python編)のQiitaにより詳しく書かれている。

本記事のまとめ

• 本記事のスタートはそもそも動的計画法の実装方法が思いつかなかったことから始まっているが、「一段階前の情報からある段階に情報を集約する」系は動的計画法による実装を疑うべし。
• 動的計画法として実装するためにうまく軸に情報集約することも必要。
  • 今回であればアルファベットに対応する軸を$\{ -, A, B, C,...,Z\}$ではなく$\{\{-\},\{-,A\},\{-,A,B\},...,\{-,A,B,...,Z\}\}$とみなすことが重要だった。
• 大きな素数で割った余りを出力する系の問題で割り算が出てきた場合は逆元を使う。
• 逆元はユーグリッド互除法を応用して高速に求めることができる。