Official

E - Mod i Editorial by penguinman


\(\text{dp}[i][j]:=(A\) の部分列 \(A'=(A_1,A_2,\ldots,A_i)\)\(j\) 個の空でない連続した部分列に切り分ける通り数\()\) と定義した動的計画法を考えます。

遷移は以下の通りです。

  • \(\text{dp}[i][j]:=(k \lt i\) を満たしてかつ \(\sum_{l=k+1}^{i}A_l\)\(j\) で割り切れるような全ての \(k\) における \(\text{dp}[k][j-1]\) の総和\()\)

上記の遷移は \(\text{sum}[0]=0,\text{sum}[i]=\text{sum}[i-1]+A_i\) を満たす配列 \(\text{sum}\) を用いて表すことも可能で、その場合以下のようになります。

  • \(\text{dp}[i][j]:=(k \lt i\) を満たしてかつ \(\text{sum}[i]-\text{sum}[k]\)\(j\) で割り切れるような全ての \(k\) における \(\text{dp}[k][j-1]\) の総和\()\)

このような dp は愚直に行うと遷移に \(O(N)\) かかることから合計での計算量が \(O(N^3)\) となり、実行時間制限に間に合いません。

そこで、「\(\text{sum}[i]-\text{sum}[k]\)\(j\) で割り切れること」と「\(\text{sum}[i]\)\(j\) で割ったあまりと \(\text{sum}[k]\)\(j\) で割ったあまりが等しいこと」が同値であることを利用して、各 \(i,j\ (0 \leq j \lt i \leq N)\) について「\(\text{sum}[k]\)\(i\) で割ったあまりが \(j\) となるような全ての \(k\) における、\(\text{dp}[k][i-1]\) の総和」をメモしながら遷移してみましょう。

これにより遷移にかかる計算量が \(O(1)\) に落ち、合計での計算量が \(O(N^2)\) となるため実行時間制限に間に合わせることができます。

実装例 (Python)

N = int(input())
A = list(map(int,input().split()))
sum_ = [0]*(N+1)
for i in range(N):
    sum_[i+1] = sum_[i]+A[i]
mem = [[0]*(N+1) for i in range(N+1)]
dp = [[0]*(N+1) for i in range(N+1)]
mem[1][0] = dp[0][0] = 1
mod = int(1e9+7)
for i in range(N):
    for j in range(1,N+1):
        dp[i+1][j] = mem[j][sum_[i+1]%j]
    for j in range(2,N+1):
        mem[j][sum_[i+1]%j] += dp[i+1][j-1]
        mem[j][sum_[i+1]%j] %= mod
ans = sum(dp[N])%mod
print(ans)

以下のように、dp 配列を使い回すような実装をすることも可能です。

実装例 (C++)

#include<bits/stdc++.h>
using namespace std;

using ll = long long;
const int mod = 1e9+7;

int main(){
    int N; cin >> N;
    vector<ll> A(N), sum(N+1);
    for(int i=0; i<N; i++){
        cin >> A[i];
        sum[i+1] = sum[i]+A[i];
    }
    vector<vector<int>> dp(N+2,vector<int>(N+2));
    dp[1][0]++;
    int ans = 0;
    for(int i=0; i<N; i++){
        for(int j=N; 1<=j; j--){
            dp[j+1][sum[i+1]%(j+1)] += dp[j][sum[i+1]%j];
            dp[j+1][sum[i+1]%(j+1)] %= mod;
            if(i == N-1){
                ans += dp[j][sum[i+1]%j];
                ans %= mod;
            }
        }   
    }
    cout << ans << endl;
}

posted:
last update: