思路

一个连续子序列（片段）可以由其起始位置 $i$ 和结束位置 $j$ 定义（其中 $1 \le i \le j \le n$）。 根据题意，一个从索引 $i$ 到 $j$ 的片段是“平衡片段”的条件是： $\sum_{k=i}^{j} A_k = j - i + 1$ 这里的 $\sum_{k=i}^{j} A_k$ 是片段中所有水晶的能量和，而 $j - i + 1$ 是该片段的长度。

直接暴力枚举所有的起始位置 $i$ 和结束位置 $j$，然后计算和并判断，时间复杂度会达到 $O(n^2)$，对于 $n=2 \times 10^5$ 的数据规模来说太慢了，会超时。我们需要一个更高效的算法。

我们可以对这个条件进行数学变换。等式右边的长度 $j - i + 1$ 可以看作是 $j - i + 1$ 个 $1$ 相加的结果。 $\sum_{k=i}^{j} 1 = j - i + 1$ 所以原条件可以写成： $\sum_{k=i}^{j} A_k = \sum_{k=i}^{j} 1$ 将等式右边移到左边： $\sum_{k=i}^{j} A_k - \sum_{k=i}^{j} 1 = 0$ 合并求和符号： $\sum_{k=i}^{j} (A_k - 1) = 0$ 这个转换是解决问题的关键。它将原问题转化为了一个更经典的问题：寻找一个新序列中，有多少个连续子序列的和为 0。

我们可以先创建一个新序列 $B$，其中 $B_k = A_k - 1$。现在问题就变成了，在序列 $B$ 中有多少个连续子片段的和为 0。

这个问题可以用前缀和 + 哈希表的方法在 $O(n)$ 的时间复杂度内解决。 我们定义前缀和数组 $P$，其中 $P_j = \sum_{k=1}^{j} B_k$。特别地，我们定义 $P_0 = 0$。 一个子序列 $B[i \dots j]$ 的和可以表示为 $\sum_{k=i}^{j} B_k = P_j - P_{i-1}$。 因此，我们要寻找的就是满足下面条件的 $(i, j)$ 对的数量： $P_j - P_{i-1} = 0$ 即： $P_j = P_{i-1}$ 所以，问题最终转化为：计算有多少对索引 $(i-1, j)$（其中 $0 \le i-1 < j \le n$）使得它们的前缀和相等。

我们可以使用一个哈希表（在 C++ 中是 unordered_map，Python 中是 dict，Java 中是 HashMap）来记录到目前为止每个前缀和值出现的次数。

算法步骤如下：

1. 初始化一个哈希表 counts，用来存储每个前缀和值出现的次数。
2. 初始化一个变量 prefix_sum 为 $0$，表示当前的前缀和。
3. 初始化结果 ans 为 $0$。
4. 在哈希表中放入 {0: 1}，即 counts[0] = 1。这代表了 $P_0=0$ 的情况，使得那些从数组开头就开始且和为 $0$ 的子数组能够被正确计数。
5. 遍历数组 $A$ 中的每个元素 $A_k$（从 $k=1$ 到 $n$）： a. 计算 $B_k = A_k - 1$。 b. 更新当前前缀和：prefix_sum = prefix_sum + B_k。 c. 在哈希表中查找 prefix_sum 这个值在之前出现过多少次，假设是 c 次。这意味着有 c 个之前的索引 $i-1$ 满足 $P_{i-1} = P_k$。这 c 个位置都可以作为起点，与当前位置 $k$ 形成和为 $0$ 的子序列。因此，我们将 ans 加上 c。 d. 将当前这个 prefix_sum 的出现次数在哈希表中加一，以供后续的计算使用。
6. 遍历结束后，ans 的值就是最终答案。

由于前缀和的值可能很大或很小，需要使用 64 位整型（long long）来存储。

C++

#include <iostream>
#include <vector>
#include <unordered_map>

int main() {
    // 提高C++的I/O效率
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);

    int n;
    std::cin >> n;

    // 使用unordered_map作为哈希表，存储前缀和及其出现的次数
    // key为前缀和的值，value为该值出现的次数
    // 使用long long防止前缀和溢出
    std::unordered_map<long long, int> prefix_sum_counts;
    
    // 初始化前缀和为0的情况
    // P_0 = 0，为了处理从数组开头就满足条件的子段
    prefix_sum_counts[0] = 1;

    long long current_prefix_sum = 0;
    long long answer = 0;

    for (int i = 0; i < n; ++i) {
        long long a;
        std::cin >> a;

        // 根据推导的公式 B_k = A_k - 1，更新前缀和
        current_prefix_sum += (a - 1);

        // 如果当前的前缀和在哈希表中已存在
        // 说明之前有一个点j，使得P_j = current_prefix_sum
        // 那么从j+1到当前位置i的子段(B[j+1]...B[i])的和就是0
        // 这个前缀和出现过多少次，就意味着有多少个这样的子段
        if (prefix_sum_counts.count(current_prefix_sum)) {
            answer += prefix_sum_counts[current_prefix_sum];
        }

        // 将当前的前缀和存入哈希表，或将其出现次数加一
        prefix_sum_counts[current_prefix_sum]++;
    }

    std::cout << answer << std::endl;

    return 0;
}

Python

import sys

def solve():
    """
    主解决函数
    """
    try:
        # 从标准输入读取n
        n = int(sys.stdin.readline())
        # 读取n个整数并存入列表
        a = list(map(int, sys.stdin.readline().split()))
    except (IOError, ValueError):
        # 处理可能的输入错误
        return

    # 使用字典作为哈希表，存储前缀和及其出现的次数
    # key为前缀和的值，value为该值出现的次数
    prefix_sum_counts = {}
    
    # 初始化前缀和为0的情况
    # P_0 = 0，为了处理从数组开头就满足条件的子段
    prefix_sum_counts[0] = 1

    current_prefix_sum = 0
    answer = 0

    for val in a:
        # 根据推导的公式 B_k = A_k - 1，更新前缀和
        current_prefix_sum += (val - 1)

        # 在字典中查找当前前缀和出现的次数
        # dict.get(key, 0) 表示如果key不存在，则返回默认值0
        # 这意味着有 `count` 个之前的子数组可以和当前位置形成满足条件的片段
        count = prefix_sum_counts.get(current_prefix_sum, 0)
        answer += count

        # 将当前的前缀和存入字典，或将其出现次数加一
        prefix_sum_counts[current_prefix_sum] = prefix_sum_counts.get(current_prefix_sum, 0) + 1

    print(answer)

if __name__ == "__main__":
    solve()

Java

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.IOException;
import java.util.HashMap;
import java.util.StringTokenizer;

public class Main {
    public static void main(String[] args) throws IOException {
        // 使用BufferedReader以提高读取效率
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        
        // 读取水晶数量n
        int n = Integer.parseInt(br.readLine());

        // 读取一行能量值
        StringTokenizer st = new StringTokenizer(br.readLine());

        // 使用HashMap作为哈希表，存储前缀和及其出现的次数
        // Key为前缀和的值，Value为该值出现的次数
        // 使用Long来防止前缀和溢出
        HashMap<Long, Integer> prefixSumCounts = new HashMap<>();
        
        // 初始化前缀和为0的情况
        // P_0 = 0，为了处理从数组开头就满足条件的子段
        prefixSumCounts.put(0L, 1);

        long currentPrefixSum = 0L;
        long answer = 0L;

        for (int i = 0; i < n; i++) {
            // 读取下一个能量值
            long a = Long.parseLong(st.nextToken());

            // 根据推导的公式 B_k = A_k - 1，更新前缀和
            currentPrefixSum += (a - 1);

            // 在哈希表中查找当前前缀和出现的次数
            // map.getOrDefault(key, 0) 表示如果key不存在，则返回默认值0
            // 这意味着有 `count` 个之前的子数组可以和当前位置形成满足条件的片段
            int count = prefixSumCounts.getOrDefault(currentPrefixSum, 0);
            answer += count;

            // 将当前的前缀和存入哈希表，或将其出现次数加一
            prefixSumCounts.put(currentPrefixSum, count + 1);
        }

        System.out.println(answer);
    }
}

题目内容

多多最近来到了一片神秘的山谷，山谷中散落着一些能量水晶。每一颗水晶拥有一个整数能量值，可能为正，也可能为负。

水晶被排成一条直线，多多决定在其中选择一个连续的水晶序列进行采集。为了保证能量的平衡，他只愿意采集“平衡片段"——即片段中所有水晶的能量值加起来，恰好等于这个片段中水晶的数量。

请你帮多多计算一下，在这条水晶链中，总共有多少个这样的“平衡片段”。

输入描述

第一行包括一个整数 $n(1 ≤n≤2×10^5)$ ，代表山谷中能量水晶的数量。

第二行包括 $n$ 个整数 $A_1, A_2,…, A_n(-10^9≤A_i≤10^9)$ ，代表第 $i$ 个水晶的能量值。

输出描述

输出一个整数，表示一共有多少个“平衡片段” 。

样例1

输入

3
2 0 1

输出

3

说明

样例中，满足条件的子段共有 $3$ 个：

$[1]$ ：能量和为 $1$ ，长度为 $1$

$[2,0]$ ：能量和为 $2$ ，长度为 $2$

$[1,2,0]$ ：能量和为 $3$ ，长度为 $3$