每日一题之 hiho233周 数组分拆

描述
小Ho得到了一个数组作为他的新年礼物，他非常喜欢这个数组！

在仔细研究了几天之后，小Ho成功的将这个数组拆成了若干段，并且每段的和都不为0！

现在小Ho希望知道，这样的拆分方法一共有多少种？

两种拆分方法被视作不同，当且仅当数组断开的所有位置组成的集合不同。

输入
每组输入的第一行为一个正整数N，表示这个数组的长度

第二行为N个整数A1~AN，描述小Ho收到的这个数组

对于40%的数据，满足1<=N<=10

对于100%的数据，满足1<=N<=105, |Ai|<=100

输出
对于每组输入，输出一行Ans，表示拆分方案的数量除以(109+7)的余数。

样例输入
5
1 -1 0 2 -2
样例输出
5

思路：

这道题是一道DP题。首先我们比较容易能想到一个O(N^2)的DP。

我们用f[i]表示将A1 ~ Ai拆分成若干段使得每段和都不为0的方案数。考虑拆出的最后一段是A[j+1] … Ai，有转移方程：

f[0] = 0

f[i] = Σf[j] | j = 0 … i-1 且 A[j+1] … Ai的和不为0。
其中“A[j+1] … Ai的和不为0”可以利用前缀和做到O(1)的判断。所以上述的DP是O(N)状态和O(N)的转移，总复杂度是O(N^2)的。

下面我们介绍如何将转移优化为O(1)的。如果我们令前缀和s[i] = f[0] + f[1] + … f[i]，那么显然上述转移方程直观上是：

f[i] 等于 f[0] … f[i-1] 中若干项的和。具体是哪些项呢？ 是那些满足s[j] != s[i]的f[j]。

于是我们可以换一种方法计算f[i]：先加上f[0] … f[i-1]，再减去那些满足s[j] == s[i]的f[j]：

f[i] = (Σf[j] | j = 0 … i-1) - (Σf[j] | j = 0 … i-1 且 s[i] == s[j])
其中f[0] + f[1] + f[2] + … + f[i-1]就是s[i-1]可以直接O(1)得到。

而第二部分“满足s[j] == s[i]的f[j]之和”我们可以用一个哈希表( unordered_map) h来记录，h[x]表示满足前缀和s[j]==x的f[j]之和。

h随着新的f[i]计算出来会更新，更新和查找的复杂度都是O(1)的。

于是我们就得到了转移O(1)的DP，总复杂度是O(N)的。

#include <bits/stdc++.h>

using namespace std;

const int mod = 1e9+7;
const int maxn = 1e5+7;

int A[maxn];
int sum[maxn];
int f[maxn];

void solve(int n) {

	f[0] = 0;
	sum[0] = 0;
	unordered_map<int,int>mp;
	for (int i = 1; i <= n; ++i) {
		sum[i] = sum[i-1] + A[i];
	}
	mp[0] = 1;
	int pre = 1;
	for (int i = 1; i <= n; ++i) {
		f[i] = (pre - mp[sum[i]] + mod) % mod;
		mp[sum[i]] = (mp[sum[i]] + f[i]) % mod;
		pre = (pre + f[i]) % mod;
		
	}
	cout << f[n] << endl;
}

int main() {

	int n;
	cin >> n;
	for (int i = 1; i <= n; ++i)
		cin >> A[i];

	solve(n);


	return 0;
}