数组拆分(微软2017笔试题)

给定一个长度为N的数组A=[A1,A2, … AN]，请你将A拆分成3段连续的子数组：
A1, A2, … Ap | Ap+1, Ap+2, … Aq | Aq+1,Aq+2, … AN

令S1 = A1 + A2 + … Ap

S2 = Ap+1 + Ap+2 + … Aq

S3 = Aq+1 + Aq+2 + … AN

问有多少种划分方案可以使得S1, S2, S3两两相差不超过1.

N <= 100000, -1000000 <= Ai <= 1000000

暴力解法：
复杂度O(n^2)

#include <iostream>
#include <cmath>
/*
测试样例 
5
3 2 1 0 2
3

6
3 2 1 0 2 1
2

*/ 
using namespace std;
const int maxn = 100005;
typedef long long ll;
int s[maxn];
int main()
{
	int n;
	cin >> n;
	for (int i = 1; i <= n; i++)
	{
		int t;
		cin >> t;
		s[i] = s[i - 1] + t;
	} 
	ll ans = 0;
	ll s1, s2, s3;
	for (int i = 1; i <= n - 2; i++)
	{
		s1 = s[i] - s[0];
		for (int j = n; j >= i + 2; j--)
		{
			s3 = s[n] - s[j - 1];
			if (abs(s3 - s1) <= 1)
			{
				s2 = s[n] - s1 - s3;
				if (abs(s1 - s2) <= 1 && abs(s2 - s3) <= 1)
				{
					ans++;
				}
			}
		}
	}
	cout << ans << endl;
	return 0;
} 

下面使用哈希优化：

#include <iostream>
#include <map>
#include <cmath>
using namespace std;
const int maxn = 100005;
typedef long long ll;
ll sum[maxn];
int a[maxn];
map<ll, int> mmp;
int main()
{
	int n;
	cin >> n;
	for (int i = 1; i <= n; i++)
	{
		cin >> a[i];
	}
	for (int i = 1; i <= n; i++)
	{
		sum[i] = sum[i-1] + a[i];//求前缀和 
		if (!mmp.count(sum[i]))//和不存在，置为0 
		{
			mmp[sum[i]] = 0;		
		}
		if(i < n)
		{
			mmp[sum[i]]++;
		}
	}
	ll ans = 0;
	ll s3 = 0;
	for(int i = n; i > 1; i--)
	{
		mmp[sum[i - 1]]--;
		s3 += a[i];
		for (ll s1 = s3 - 1; s1 <= s3 + 1; s1++)// 不妨假设|s1-s3| <= 1，枚举s1 
		{
			ll s2 = sum[n] - s1 - s3;
			if (abs(s1 - s2) <= 1 && abs(s2 - s3) <= 1)
			{
				ans += mmp[s1];//s1这个和存在与否不影响，因为不存在mmp[s1] = 0 
			}
		}
	}
	cout << ans << endl; 
	return 0;
}