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: