【前缀和】截断数组

题目

3956. 截断数组 - AcWing题库

---

给定一个长度为 $n$ 的数组 $a_1,a_2,…,a_n$ 。

现在，要将该数组从中间截断，得到三个非空子数组。

要求，三个子数组内各元素之和都相等。

请问，共有多少种不同的截断方法？

输入格式

第一行包含整数 $n$ 。

第二行包含 $n$ 个整数 $a_1,a_2,…,a_n$ 。

输出格式

输出一个整数，表示截断方法数量。

数据范围

前六个测试点满足 $1 \le n \le 10$ 。
所有测试点满足 $1 \le n \le 10^5$ ， $−10000 \le a_i \le 10000$ 。

输入样例1：

4
1 2 3 3

输出样例1：

1

输入样例2：

5
1 2 3 4 5

输出样例2：

0

输入样例3：

2
0 0

输出样例3：

0

解题

方法一：前缀和 枚举

思路

由【转载】由数据范围反推算法复杂度以及算法内容可以看出这题的时间复杂度要控制在 $O(n \log n)$ 之内，那么也就不能采取暴力枚举两个截断点这样 $O(n^2)$ 的做法了。

首先，如果数组的长度小于 $3$ 或数组所有元素和不能被 $3$ 整除时一定不存在截断方法，直接输出 $0$ 。然后可以把数组预处理为前缀和数组以快速的判断子数组和。再然后枚举第一个截断点，并处理为前缀和数组，这样就可以在最后枚举第二个截断点时以 $O(1)$ 的时间复杂度快速地查询有多少个第一截断点可以与其搭配为截断方案。

例如（以下表述数组下标均从 $1$ 开始），当 $a = [-5, 2, -1, -1, 5, 0, -4, 4, 0]$ 时：

• 首先判断其长度大于 $3$ ，且 $sum(a) \mod 3 = 0$ 说明该数组有截断方案。
• 然后把它处理为前缀和数组：$s = [-5, -3, -4, -5, 0, 0, -4, 0, 0]$ 。
• 从 $s[2 \dots n-1]$ 中枚举其第一个截断点得到截断点数组 $c = [0, 0, 0, 0, 0, 1, 1, 0, 0]$ ，这表示：
  • 在下标 $5$ 与 $6$ 之间可以存在一个截断点使 $a$ 被截断为：$[-5, 2, -1, -1, 5 \,\, | \,\, 0, -4, 4, 0]$ 。
  • 在下标 $6$ 与 $7$ 之间可以存在一个截断点使 $a$ 被截断为：$[-5, 2, -1, -1, 5, 0 \,\, | \,\, -4, 4, 0]$ 。
• 把 $c$ 处理为前缀和数组：$cs = [0, 0, 0, 0, 0, 1, 2, 2, 0]$ 。
• 从 $s[n \dots 3]$ 中枚举其第一个截断点：
  • 当 $i=9$ 时，$s_n - s_{i-1} = 0$ ，说明在下标 $8$ 与 $9$ 之间可以存在一个截断点使 $a$ 被截断为：$[-5, 2, -1, -1, 5, 0, -4, 4 \,\, | \,\, 0]$ ，查截断点数量前缀和数组 $cs_{i-1} = 2$ 得第二截断点选为此处时有 $2$ 个第一截断点可供选择，故截断方法数量增加 $2$ （cnt += cs[i - 1]）。
  • …
  • 当 $i=7$ 时，$s_n - s_{i-1} = 0$ ，说明在下标 $6$ 与 $7$ 之间可以存在一个截断点使 $a$ 被截断为：$[-5, 2, -1, -1, 5, 0 \,\, | \,\, -4, 4, 0]$ ，查截断点数量前缀和数组 $cs_{i-1} = 1$ 得第二截断点选为此处时有 $1$ 个第一截断点可供选择，故截断方法数量增加 $1$ （cnt += cs[i - 1]）。
  • …

注意：数组长度为 $10^5$ 时，截断方法数量最大可能到 $\sum_{i=1}^{99998} i = \frac{99999 * 99998}{2}= 4999850001$ 造成整数溢出，所以 cnt 要用长整形保存。

代码

import java.util.*;
import java.io.*;

public class Main {
    public static void main(String[] args) throws IOException {
        StreamTokenizer in = new StreamTokenizer(new BufferedReader(new InputStreamReader(System.in)));
        in.nextToken();
        int n = (int) in.nval;
        if (n < 3) {
            System.out.println(0);
            return;
        }
        int[] s = new int[n + 1];
        for (int i = 1; i <= n; ++i) {
            in.nextToken();
            s[i] = s[i - 1] + (int) in.nval;
        }
        if (s[n] % 3 != 0) {
            System.out.println(0);
            return;
        }
        int avg = s[n] / 3;
        int[] cs = new int[n + 1];
        for (int i = 2; i <= n - 1; ++i) {
            if (s[i - 1] == avg) cs[i] = 1;
            cs[i] += cs[i - 1];
        }
        long cnt = 0;
        for (int i = n; i >= 3; --i) {
            if (s[n] - s[i - 1] == avg) {
                cnt += cs[i - 1];
            }
        }
        System.out.println(cnt);
    }
}