Official

E - Mod i Editorial by en_translator


Let us consider a DP (Dynamic Programming) where we define \(\text{dp}[i][j]:=(\)the number of partitions of \(A'=(A_1,A_2,\ldots,A_i)\) into \(j\) non-empty contiguous subsequences\()\).

The transitions are as follows:

  • \(\text{dp}[i][j]:=(\) The sum of \(\text{dp}[k][j-1]\) for all \(k \lt i\) such that \(\sum_{l=k+1}^{i}A_l\) is a multiple of \(j)\).

The transitions above can be represented by an array \(\text{sum}\) such that \(\text{sum}[0]=0,\text{sum}[i]=\text{sum}[i-1]+A_i\):

  • \(\text{dp}[i][j]:=(\) The sum of \(\text{dp}[k][j-1]\) for all \(k \lt i\) such that \(\text{sum}[i]-\text{sum}[k]\) is a multiple of \(j)\).

Such a DP costs a total of \(O(N^3)\) computations since each transition takes \(O(N)\) time, which does not fit in the time limit.

However, we have the fact that “\(\text{sum}[i]-\text{sum}[k]\) is divisible by \(j\)” if and only if “the remainder of \(\text{sum}[i]\) by \(j\) and that of \(\text{sum}[k]\) by \(j\) are equal.” What if we memorize for each \(i,j\ (0 \leq j \lt i \leq N)\) “the sum of \(\text{dp}[k][i-1]\) for all \(k\) such that the remainder of \(\text{sum}[k]\) by \(i\) is \(j\)?”

Now each transition costs only \(O(1)\) time, so that the total time complexity becomes \(O(N^2)\), which fits in the time limit.

Sample code (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)

One can reuse the DP array as follows.

Sample code (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: