[ABC262] D - I Hate Non-integer Number

[Q] https://atcoder.jp/contests/abc262/tasks/abc262_d

・考察
枝切DFSと勘違い。手元でさんざんTLE寄り道した。
DPに切り替えてAC。

1.
1個とる場合に平均値が整数になるのは何通り？
2個とる場合に平均値が整数になるのは何通り？
3個とる場合に平均値が整数になるのは何通り？
... で検証する必要がある。O(N)

2. DP[index][とった項数][その時のあまり]で管理。O(N^3)
あわせて O(N^4) で間に合いそう。
3. indexの遷移を考える。前状態を"引き継いだまま"、次のindexの状態数を加算するDPになるので、indexは省略。DPテーブルを上書き加算しようと思う。
結局 DP[とった項数][その時のあまり] で管理。俺はいつも貰うDP教徒。

・DPシミュレート

N=3
A[]: 2 6 2
を考える。

・1個とる場合
どれとってもいいので、N通り。
ans = N

・2個とる場合
1. A[]を2の剰余に変換する
2 6 2 -> 0 0 0


- DP初期状態
DP[0][0] = 1通り

2. indexを遷移していく
- 1個目の0をとる
<とる状態の加算>
DP[1][0] = 1通り // DP[0][0] -> DP[1][0]
<引継ぎ>
DP[0][0] = 1通り

<まとめて>
DP[1][0] = 1通り
DP[0][0] = 1通り


- 2個目の0をとる
<とる状態の加算>
DP[2][0] = 1通り // DP[1][0] -> DP[2][0]
DP[1][0] = 1通り // DP[0][0] -> DP[1][0]

<引継ぎ>
DP[1][0] = 1通り
DP[0][0] = 1通り

<まとめて>
DP[2][0] = 1通り
DP[1][0] = 2通り
DP[0][0] = 1通り


- 3個目の0をとる
<とる状態の加算>
DP[3][0] = 1通り
DP[2][0] = 2通り
DP[1][0] = 1通り

<引継ぎ>
DP[2][0] = 1通り
DP[1][0] = 2通り
DP[0][0] = 1通り

<まとめて>
DP[3][0] = 1通り
DP[2][0] = 3通り
DP[1][0] = 3通り
DP[0][0] = 1通り

解答は、2個選んであまり0になる状態数なので、
DP[2][0] = 3通り。
ans += 3 // 6通り




・3個とる場合
1. A[]を3の剰余に変換する
2 6 2 -> 2 0 2

- DP初期状態
DP[0][0] = 1通り

2. indexを遷移していく
- 1個目の2をとる
<とる状態の加算>
DP[1][2] = 1通り
<引継ぎ>
DP[0][0] = 1通り
<まとめて>
DP[1][2] = 1通り
DP[0][0] = 1通り


- 2個目の0をとる
<とる状態の加算>
DP[2][2] = 1通り
DP[1][0] = 1通り

<引継ぎ>
DP[1][2] = 1通り
DP[0][0] = 1通り

<まとめて>
DP[2][2] = 1通り
DP[1][0] = 1通り
DP[1][2] = 1通り
DP[0][0] = 1通り


- 3個目の2をとる
<とる状態の加算>
DP[3][1] = 1通り // DP[2][2] -> DP[3][4] = DP[3][1]
DP[2][2] = 2通り // DP[1][0] -> DP[2][2]
DP[2][1] = 1通り // DP[1][2] -> DP[2][4] = DP[2][1]
DP[1][2] = 1通り // DP[0][0] -> DP[1][2]

<引継ぎ>
DP[2][2] = 1通り
DP[1][0] = 1通り
DP[1][2] = 1通り
DP[0][0] = 1通り

<まとめて>
DP[3][1] = 1通り
DP[2][1] = 1通り
DP[2][2] = 3通り
DP[1][0] = 1通り
DP[1][2] = 2通り
DP[0][0] = 1通り

解答は、3個選んであまり0になる状態数なので、
DP[3][0] = 0通り
ans += 0 // 6

以上で
ans = 6通り

・実装

int main(){
 cincout();
 
 ll N;
 cin >> N;
 
 ll A[N];
 rep(i, N) cin >> A[i];
 mint ans = N; // 1個ずつ取ればN通り
 
 for(ll d = 2; d <= N; ++d){ // 選ぶ項数
   mint DP[d + 1][d + 1]{}; // DP[選んだ項数][あまり]
   DP[0][0] = 1; // 初期状態。DP[0個とった][あまりは0]
   rep(i, N){ // indexを動かす
     ll a = A[i] % d; // dの剰余にデータ変換
     // aをとる場合を考える
     for(ll j = min(i, d - 1); j >= 0; --j){ // 選んだ項数を全探索
       for(ll k = 0; k < d; ++k){ // あまりを全探索
         DP[j + 1][(k + a) % d] += DP[j][k];
       }
     }
   }
   // 解答を加算
   ans += DP[d][0]; // d個選んで余りが0の状態数
 }
 cout << ans << endl;
}

https://atcoder.jp/contests/abc262/submissions/33678403