Codeforces 1485F Copy or Prefix Sum

題目連結

點我跳轉

題目大意

給定一個長度為 \(N\) 的序列 \(bi\)

問有多少個長度為 \(N\) 的序列 \(a\) 使得

• \(b[i] = a[i]\)

或

• \(b[i] = ∑a[j] , j∈[1,i]\)

解題思路

定義 $dp[i][j] $ 表示前 \(i\) 項的字首和為 \(j\) 的序列 \(a\) 的個數，其中 \(dp[0][0] = 1\)

( 因為字首和很大，所以需要用 \(map\) 來操作 )

那麼：

1. 當 \(b[i] = a[i]\) 時 , \(dp[i][j] = dp[i - 1][j - b[i]]\)

2. 當 \(b[i] = ∑a[j] , j∈[1,i]\) 時 , \(dp[i][b[i]] = ∑dp[i - 1][j],j∈[-inf,inf]\)

對於第一種轉移相當於把整個陣列的值向右平移 \(b[i]\)

對於第二種轉移需要注意當 \(sum[i-1] = 0\) 時，\(b[i]\) 既等於 \(a[i]\) 又等於 \(∑a[j] , j∈[1,i]\)

相當於多轉移了一次 \(dp[i - 1][0]\) ，所以需要減去 \(dp[i - 1][0]\)

最後的答案 \(ans = ∑dp[n][j],j∈[-inf,inf]\) ，複雜度為 \(N^2logN\) ( \(log\) 的複雜度源於 \(map\) )

考慮優化：

定義 \(sum[i]\) 表示長度為 \(i\) 且滿足題目條件的序列 \(a\) 的個數

對於第一種轉移，只是把數值向右平移，並不會導致答案的個數增加，所以 \(sum[i] = sum[i - 1]\)

對於第二種轉移，\(dp[i][b[i]] += sum[i - 1]\) , 同時減去 \(dp[i - 1][0]\) ，相當於 \(sum[i] += sum[i - 1] , sum[i] -= dp[i - 1][0]\)

於是我們可以定義 \(py\) 表示平移的長度，起初 \(py = 0\)，每計算完 \(sum[i]\) 後 , 令 \(py += b[i]\)

那麼 \(dp[i - 1][j]\) 則可以用 \(dp[j - py]\) 表示

而 \(dp[i][j]\) 則可以用 \(dp[j - py - b[i]]\) 表示

於是可得 \(sum[i] += sum[i - 1] - dp[0 - py]\) , \(dp[b[i] - py - b[i]] += sum[i - ] - dp[0 - py]\)

最後的答案 \(ans = sum[n]\) , 複雜度為 \(NlogN\)

AC_Code

#include<bits/stdc++.h>

#define int long long

using namespace std;

const int N = 3e5 + 10 , mod = 1e9 + 7;

map<int , int>dp;

int b[N];

signed main()
{
	int T = 1;
	
	cin >> T;
	
	while(T --)
	{
		dp.clear();
		
		int n , sum = 1 , py = 0;
		
		cin >> n;
		
		for(int i = 1 ; i <= n ; i ++) cin >> b[i];	
		
		dp[0] = 1;
		
		for(int i = 1 ; i <= n ; i ++)
		{
			int add = (sum - dp[0 - py] + mod) % mod;
			
			sum = (sum + add) % mod , py += b[i];
			
			dp[b[i] - py] = (dp[b[i] - py] + add) % mod;
		}
		
		cout << sum << '\n';
	}
	
	return 0;
}