[USACO23FEB] Equal Sum Subarrays G 题解

题目链接

\(O(n^5)\) 暴力

显然，如果修改 \(a_i\) 的值，只会影响包含 \(a_i\) 的区间的区间和。于是对于每个 \(a_i\)，可以将所有区间分成两类，即包含 \(a_i\) 的区间和不包含 \(a_i\) 的区间。这两种区间的区间和中最小的差值即为答案。

对于每个 \(a_i\)，暴力枚举所有的区间和是 \(O(n^2)\) 的（这里可以用前缀和快速求出区间和）；枚举两类区间的差值是 \(O(n^4)\) 的——注意区间的个数是 \(O(n^2)\) 量级的，双层的遍历就是 \(O(n^4)\)。乘上外层枚举 \(a_i\) 的循环，总的时间复杂度为 \(O(n^5)\)。

部分代码如下：

// s1[]：包含 a[i] 的区间和，s2[]：不包含 a[i] 的区间和
// top1 和 top2 分别表示 s1[] 和 s2[] 当前的元素个数
for(int i = 1; i <= n; i++)
{
    top1 = top2 = 0, ans = INF;
    for(int j = 1; j <= n; j++)
    {
        for(int k = j; k <= n; k++) // 枚举 [j, k] 区间
        {
            ll s = sum[k] - sum[j-1];
            if(j <= i && k >= i) s1[++top1] = s; // 判断该区间是否包含 a[i]
            else s2[++top2] = s;
        }
    }
    for(int i = 1; i <= top1; i++) // O(n^4) 暴力枚举差值
        for(int j = 1; j <= top2; j++) ans = min(ans, abs(s1[i] - s2[j]));
    printf("%lld\n", ans);
}

\(O(n^3 \log (n^2))\)

上面这种做法中用 \(O(n^4)\) 暴力枚举差值实在太过愚蠢了。考虑优化这一步。

先将 s1[],s2[] 排序，然后使用双指针法，这样只要遍历一遍，时间复杂度降低到 \(O(n^2)\)。

这部分代码如下：

sort(s1 + 1, s1 + top1 + 1), sort(s2 + 1, s2 + top2 + 1);
int p1 = 1, p2 = 1; // p1 和 p2 分别是 s1[] 和 s2[] 的指针
s2[top2+1] = INF; // 由于双指针法可能取到 s2[top2+1]，需要先将其设为无穷大
for(p1; p1 <= top1; p1++)
{
    while(s2[p2+1] <= s1[p1] && p2 < top2) p2++;
    ans = min({ans, abs(s2[p2] - s1[p1]), abs(s2[p2+1] - s1[p1])});
    // s2[p2] 是 s2[] 中小于等于 s1[p1] 的最大的数，s2[p2+1] 是 s2[] 中大于 s1[p1] 的最小的数
    // 由于 s1[] 和 s2[] 都是升序的，所以与 s1[p1] 差值最小的数一定在这两者之中 
}

给 \(n^2\) 个数排序是 \(O(n^2 \log (n^2))\) 的，双指针法是 \(O(n^2)\)，总的时间复杂度为 \(O(n^3 \log (n^2))\)。

\(O(n^3)\)

上面的做法的时间瓶颈来自给 \(n^2\) 个数排序。如何优化呢？

我们发现 \(n\) 的范围很小，这启示我们可以把 \(O(n^2)\) 个区间和预处理，将其离散化和排序，之后每次排序 s1[],s2[] 时就可以用桶排，时间复杂度降到 \(O(n^2)\)。

具体做法详见代码：

// P9127 [USACO23FEB] Equal Sum Subarrays G
#include<bits/stdc++.h>

using namespace std;

typedef long long ll;
const int MAXN = 510;
const ll INF = 0x3f3f3f3f3f3f3f3f;
int n, top0, top1, top2;
bool vis1[MAXN * MAXN], vis2[MAXN * MAXN];
ll a[MAXN], sum[MAXN], s1[MAXN * MAXN], s2[MAXN * MAXN];
ll q[MAXN * MAXN], tot, rnk[MAXN][MAXN];

void pre()
{
	// 预处理，离散化，排序 
	for(int i = 1; i <= n; i++)
	{
		for(int j = i; j <= n; j++) q[++tot] = sum[j] - sum[i-1];
	}
	sort(q + 1, q + 1 + tot);
	top0 = unique(q + 1, q + 1 + tot) - q - 1;
	// q 是离散化后的数组 
	for(int i = 1; i <= n; i++)
	{
		for(int j = 1; j <= n; j++)
		{
			ll s = sum[j] - sum[i-1];
			rnk[i][j] = (lower_bound(q + 1, q + 1 + top0, s) - q);
			// rnk[i][j] 表示 [i,j] 的区间和的排名，即新的下标
			// q[rnk[i][j]] 即为 [i,j] 的区间和 
		}
	}	
}

ll solve(int pos)
{
	ll ans = INF;
	memset(vis1, 0, sizeof(vis1)), memset(vis2, 0, sizeof(vis2));
	top1 = top2 = 0, ans = INF;
	for(int i = 1; i <= n; i++)
	{
		for(int j = i; j <= n; j++) // [i, j]
		{
			if(i <= pos && j >= pos) vis1[rnk[i][j]] = true;
			else vis2[rnk[i][j]] = true;
		}
	}
	for(int i = 1; i < MAXN * MAXN; i++) // 桶排序
	{
		if(vis1[i]) s1[++top1] = q[i];
		else if(vis2[i]) s2[++top2] = q[i];
	}
	int p1 = 1, p2 = 1; // p1 和 p2 分别是 s1[] 和 s2[] 的指针
	s2[top2+1] = INF; // 由于双指针法可能取到 s2[top2+1]，需要先将其设为无穷大
	for(p1; p1 <= top1; p1++)
	{
	    while(s2[p2+1] <= s1[p1] && p2 < top2) p2++;
	    ans = min({ans, abs(s2[p2] - s1[p1]), abs(s2[p2+1] - s1[p1])});
	    // s2[p2] 是 s2[] 中小于等于 s1[p1] 的最大的数，s2[p2+1] 是 s2[] 中大于 s1[p1] 的最小的数
	    // 由于 s1[] 和 s2[] 都是升序的，所以与 s1[p1] 差值最小的数一定在这两者之中 
	}
	return ans;
}

int main()
{
	scanf("%d", &n);
	for(int i = 1; i <= n; i++) scanf("%lld", &a[i]), sum[i] = sum[i-1] + a[i];
	pre();
	for(int i = 1; i <= n; i++) printf("%lld\n", solve(i));
	return 0;
}

总结

这道题中，我们先思考了最暴力的做法，然后用双指针法和离散化两个技巧优化了时间复杂度。其中离散化的技巧是具有启发性的，它启示我们在数据规模较小的情况下可以先离散化，然后就可以使用一些时间复杂度更小的算法（如这道题中的桶排序）。