题解 | 喜欢切数组的红

def solve():
    # 读取输入
    n = int(input())
    arr = list(map(int, input().split()))

    # 如果数组长度小于3，无法分成3部分，直接输出0
    if n < 3:
        print(0)
        return

    # 计算数组总和
    total_sum = sum(arr)

    # 如果总和不能被3整除，无法分成3个和相等的子数组，直接输出0
    if total_sum % 3 != 0:
        print(0)
        return

    # 计算每个子数组的目标和
    target = total_sum // 3

    # 计算前缀和数组，prefix_sum[i]表示前i个元素的和
    prefix_sum = [0] * (n + 1)
    for i in range(1, n + 1):
        prefix_sum[i] = prefix_sum[i - 1] + arr[i - 1]

    # 计算前缀正数个数数组，prefix_positive[i]表示前i个元素中正数的个数
    prefix_positive = [0] * (n + 1)
    for i in range(1, n + 1):
        prefix_positive[i] = prefix_positive[i - 1] + (1 if arr[i - 1] > 0 else 0)

    # 寻找所有可能的第一分割点i，满足前i项和等于target且至少包含一个正数
    candidates_i = []
    for i in range(1, n - 1):  # i最多到n-2，确保后面至少有两个元素
        if prefix_sum[i] == target and prefix_positive[i] >= 1:
            candidates_i.append(i)

    # 寻找所有可能的第二分割点j，满足前j项和等于2*target且第三部分至少包含一个正数
    candidates_j = []
    for j in range(2, n):  # j最多到n-1，确保第三部分至少有一个元素
        if (
            prefix_sum[j] == 2 * target
            and (prefix_positive[n] - prefix_positive[j]) >= 1
        ):
            candidates_j.append(j)

    # 统计满足条件的分割方案数
    count = 0
    for j in candidates_j:
        # 遍历所有i < j的候选分割点i，确保中间部分至少有一个正数
        for i in candidates_i:
            if i >= j:
                continue  # 确保i < j
            if prefix_positive[j] - prefix_positive[i] >= 1:
                count += 1

    print(count)


solve()