ABC262 D - I Hate Non-Integer Number Python解答例

AtCoder Beginner Contest 262 D - I Hate Non-Integer NumberをPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

項数が$N$の正整数列$A=(a_1,\ldots,a_N)$が与えられます。
$A$の項を$1$個以上選ぶ方法は$2^N-1$通りありますが、そのうち選んだ項の平均値が整数であるものが何通りかを$998244353$で割った余りを求めてください。

制約

• $1 \leq N \leq 100$
• $1 \leq a_i \leq 10^9$
• 入力はすべて整数

解答例

動的計画法

$A$の項から$k$個選んだときの平均値が整数であるためには、選んだ数の総和が$k$で割り切れる必要があります。

ここで、関心があるのは$k$個選んだときの総和を$k$で割った余りであり、選ぶ項の数は$N$個を超えることはありません。
また、$i - 1$個選んだときの総和を$k$で割った余りがわかっているとき、追加で1つ選んだときの余りは選んだ数を足すことにより計算することができます。

なんとなく動的計画法で解ける気がします。

DPテーブルの定義

整数$l (1 \leq l \leq N)$のそれぞれに対して、DPテーブルを以下のように定義します。

$dp[i][j][k]$：$A$の先頭から$i$個を選択対象として、$j$個の項を選んだとき、選んだ数の総和を$l$で割った余りが$k$になるような選び方の場合の数

ただし、$1 \leq i \leq N, 1 \leq j \leq l, 0 \leq k < l$です。

初期条件は$dp[i][0][0] = 1$です。1個も選ばなければ総和も余りも0です。

遷移は、$i$番目の項$a_i$を選ぶか選ばないかで分岐します。

• $a_i$を選ばない場合：$dp[i][j][k] += dp[i - 1][j][k]$
• $a_i$を選ぶ場合：$dp[i][j][k] += dp[i - 1][j - 1][k - a_i \mod l] \quad (k - a_i \geq 0)$

この問題の解は、全ての$l$について$dp[N][l][0]$を足し合わせた数になります。

遷移の解釈

• $a_i$を選ばない場合、$i - 1$項までを選択対象にしたときの結果から変動しないので、結果をそのまま足す。
• $a_i$を選ぶ場合、「$a_i$を選んだことにより、選んだ項の個数数が$j$になり、総和を$l$で割った余りが$k$になる」場合の数は、$dp[i - 1][j - 1][k - a_i]$の場合の数に等しいため、これを足す。

DPテーブルの定義について補足

この解法では、1以上$N$以下の整数$l$のそれぞれについて個別に定義し、計算した結果を集約して解を求めています。

「以下のようなDPテーブル1つで解が求まるんじゃないか」と一瞬考えますが、それはできないことがすぐにわかります。

$dp[i][j][k]$: $A$の先頭から$i$個を選択対象として、$j$個の項を選んだとき、選んだ数の総和を$j$で割った余りが$k$になるような場合の数

このDPテーブルで遷移を考えたとき、$i$番目の項$a_i$を選んだ場合で、選んだ項の個数が$j$になります。
このときの場合の数は、選んだ項の個数が$j - 1$個であったときの結果($dp[i - 1][j - 1][k]$)から求めなければなりませんが、これらの結果が持つ$k$の値は、$j - 1$で割ったときの余りです。そのため、それまでに選んだ項の総和の情報を得ることができず、$j$で割ったときの余りを求めることができません。

なので、選んだ数の総和を割る数を$l$で固定し、各$l$について動的計画法を適用して解を求める必要があります。

($N$と$a_i$の値が小さかったら1つのDPテーブルだけで解けるのかな？)

実装例

Pythonによる実装例を以下に示します。PyPy3で提出しないとTLEします。
$k - a$を$l$で割った余りは、Pythonならば(k - a % l) % lと書くことができます。

d.py

from typing import List


def main():
    # 入力受け取り
    N: int = int(input())
    A: List[int] = list(map(int, input().split()))

    mod: int = 998244353

    answer: int = 0
    # 1以上N以下の整数lについてそれぞれ動的計画法で解を求める
    for l in range(1, N + 1):
        # DPテーブル
        # 内側の次元のリストの長さを制限しないとTLEするので注意
        dp: List[List[List[int]]] = [
            [[0 for k in range(l)] for j in range(l + 1)] for i in range(N + 1)
        ]

        # 初期条件
        for i in range(N + 1):
            dp[i][0][0] = 1

        # 遷移
        for i in range(1, N + 1):
            for j in range(1, l + 1):
                for k in range(l):
                    # i番目の要素を選ばない場合
                    dp[i][j][k] += dp[i - 1][j][k]

                    # i番目の要素を選ぶ場合
                    dp[i][j][k] += dp[i - 1][j - 1][(k - (A[i - 1] % l)) % l]

                    # 余りを取る
                    dp[i][j][k] %= mod

        # 解を足し合わせる
        answer += dp[N][l][0] % mod

    print(answer % mod)


if __name__ == "__main__":
    main()

実際の提出はこちら。
(PyPy3でも実行制限時間ギリギリなのでもうちょっと改善の余地がありそうです。)