[数据结构与算法]leetcode 560.和为K的子数组 前缀和+哈希表优化

题目描述

和为K的子数组
给你一个整数数组 $nums$ 和一个整数 $k$ ，请你统计并返回 该数组中和为 $k$ 的连续子数组的个数 。

思路

考虑以 $i$ 结尾和为 $k$ 的连续子数组个数，我们需要统计符合条件的下标 $j$ 的个数，其中 $0\le j\le i$ 且$[j..i]$ 这个子数组的和恰好为 $k$ 。

所以我们考虑以 $i$ 结尾的和为 $k$ 的连续子数组个数时只要统计有多少个前缀和为 $pre[i]?k$ 的 $pre[j]$ 即可。我们建立哈希表 $mp$，以和为键，出现次数为对应的值，记录$pre[i]$ 出现的次数，从左往右边更新 $mp$ 边计算答案，那么以 $i$ 结尾的答案 $mp[pre[i]?k]$ 即可在 $O(1)$ 时间内得到。最后的答案即为所有下标结尾的和为 $k$ 的子数组个数之和。

需要注意的是，从左往右边更新边计算的时候已经保证了$mp[pre[i]?k]$ 里记录的 $pre[j]$ 的下标范围是 $0\le j\le i$ 。同时，由于$pre[i]$ 的计算只与前一项的答案有关，因此我们可以不用建立 $pre$ 数组，直接用$pre$ 变量来记录 $pre[i?1]$ 的答案即可。

代码

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        unordered_map<int,int> mp;  //哈希表

        int count=0;  //出现次数初始化为0
        mp[0]=1;  //特殊情况，从0加到i的和为k，j为0
        int pre=0;  //前缀和初始为0

        for(auto &x:nums)
        {
            pre+=x;  //前缀和加上来
            if(mp.find(pre-k)!=mp.end())  //如果出现过
                count+=mp[pre-k];    //加上出现的次数
            mp[pre]++;  //当前前缀和出现了
        }
        return count;
    }
};