题目介绍

本题为“门门CS 2008”题库的第 31 题，题目描述如下：

有一个长度为 $n$ 的数组 $a_1,a_2,\dots,a_n$，请你计算出其中有多少个子数组，满足其元素之和恰好为 $k$。

具体的输入和输出格式请参考下文。

解题思路

要计算出子数组元素之和恰好为 $k$ 的数量，可以使用前缀和来计算。具体来说，对于数组 $a$，可以维护一个前缀和数组 $s$，其中 $s_i=\sum_{j=1}^i a_j$。那么对于区间 $[l,r]$ 的元素之和就可以通过计算 $s_r-s_{l-1}$ 来得到。

接下来考虑如何统计子数组个数。假设当前处理到位置 $i$，那么可以在前面的 $i-1$ 个位置中查找有多少个位置 $j$ 满足 $s_i-s_j=k$。具体来说，可以使用一个哈希表来维护前缀和出现的次数。假设当前处理到位置 $i$，那么可以在哈希表中查找 $s_i-k$ 出现的次数，即为以位置 $i$ 结尾的子数组中元素之和为 $k$ 的个数。将这些个数累加起来即可得到总的子数组个数。

输入格式

第一行包含两个整数 $n$ 和 $k$，表示数组 $a$ 的长度和目标和。

第二行包含 $n$ 个用空格隔开的整数 $a_1,a_2,\dots,a_n$，表示数组 $a$。

输出格式

输出一个整数，表示满足条件的子数组数量。

输入样例

4 0
1 -1 2 -2

输出样例

5

参考代码

def subarray_sum(nums, k):
    s = res = 0
    cnt = {0: 1}
    for num in nums:
        s += num
        res += cnt.get(s - k, 0)
        cnt[s] = cnt.get(s, 0) + 1
    return res

n, k = map(int, input().split())
a = list(map(int, input().split()))

print(subarray_sum(a, k))

subarray_sum 函数即为本题的主要解题函数，输入 $nums$ 表示数组 $a$，$k$ 表示目标和，返回满足条件的子数组数量。在函数中，$s$ 表示当前处理到的前缀和，$res$ 表示满足条件的子数组数量，$cnt$ 是一个哈希表，存储前缀和出现的次数。对于数组 $nums$ 中的每个元素 $num$，先更新前缀和 $s$，然后在哈希表 $cnt$ 中查找 $s-k$ 出现的次数，更新 $res$，最后在哈希表 $cnt$ 中添加新的 $s$，并将其出现的次数加 1。最后返回结果 $res$ 即可。