【LeetCode专题】前缀和+哈希表

LeetCode 560. 和为 K 的子数组

题目 示例

输入：nums = [1,1,1], k = 2
输出：2

输入：nums = [1,2,3], k = 3
输出：2

思路：

利用前缀和sum记录当前的和，区间[l,r]的和为: s[r+1] - s[l]
nums    1 0 1 1 0 2  [k = 2]
sum   0 1 1 2 3 3 5

利用哈希表map来存储 [历史前缀和], 题目要找的是 [当前前缀和 - 历史前缀和 == k] 的次数
所以需要遍历当前前缀和，如果 [sum - k] 出现在map中 cnt 次， 则结果 + cnt次，有 cnt 个子数组满足题意。
例如 i = 4时 sum = 3, sum - k = 1, mp[1] = 2, 说明有2个子数组满足, 即 [0,1,1]、[1,1]

代码：

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        int n = nums.size(), cnt = 0;
        int sum = 0;
        // map优化
        unordered_map<int, int> mp;
        // for(int i = 0; i < n; i++) s[i+1] = s[i] + nums[i]; // 必须遍历过程中计算前缀和
        mp[0] = 1; // 处理边界：前缀和为 0 出现了 1 次
        for(int r = 0; r < n; r++){
            sum += nums[r];
            if(mp.count(sum - k) > 0){
                cnt += mp[sum - k];
            }
            mp[sum] ++;
        }
        return cnt;
    }
};

LeetCode 525. 连续数组

题目 示例

输入: nums = [0,1,0]
输出: 2
说明: [0, 1] (或 [1, 0]) 是具有相同数量0和1的最长连续子数组。

思路

1 表示 +1
0 表示 -1
用前缀和来表示1和0的数量
同样, 用map来存储历史前缀和第一次出现时的[下标]
nums    0 1  0  0  1 1
sum  0 -1 0 -1 -2 -1 0
遍历前缀和，如果当前前缀和在历史前缀和中出现过，说明区间 [l,r] 具有相同的0和1，更新ans
否则，记录当前前缀和第一次出现时的下标pos

代码

class Solution {
public:
    int findMaxLength(vector<int>& nums) {
        int ans = 0, n = nums.size();
        unordered_map<int, int> mp;
        mp.insert({0,-1}); //初始 0 的 pos = -1
        int sum = 0; 
        for(int r = 0; r < n; r++){
            // 一边遍历一边计算前缀和
            sum += nums[r] == 1 ? 1 : -1; // core/
            if(mp.count(sum) > 0){
                ans = max(ans, r - mp[sum]);
            }
            else mp[sum] = r;
        }
        return ans;
    }
};

LeetCode 974. 和可被 K 整除的子数组

题目 示例

输入：nums = [4,5,0,-2,-3,1], k = 5
输出：7
解释：
有 7 个子数组满足其元素之和可被 k = 5 整除：
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]

思路

[l,r]区间和: s[r+1] - s[l]
(s[r+1] - s[l]) % k == 0
<==> s[r+1] % k - s[l] % k == 0
<==> s[r+1] % k == s[l] % k
所以区间和 mod k == 0 的个数  等价于 前缀和s[r+1] % k 等于 s[l] % K 的个数
例如：
nums    4  5 0 -2 -3 1     k = 5
sum   0 4  9 9  7  4 5
%k    0 1  1  1  2  1 0
mp[1] = 4, mp[0] =2
前缀和%k = 1的次数为4， %k = 0的次数为2
cnt = ∑ N * (N-1) / 2 = 6 + 1 = 7

代码

class Solution {
public:
    int subarraysDivByK(vector<int>& nums, int k) {
        // num   4 5 0 -2 -3 1
        // s   0 4 9 9  7  4 5
        //%k   0 1 1 1  2  1 0
        // [l, r] -->  s[r+1] - s[l]
        // s[r+1] - s[l] % k == 0 ?
        // s[r+1] % k == s[l] % k ?
        // mp[1] = 4 mp[0]=2 
        // n*(n-1) /2 
        int n = nums.size();
        long long cnt = 0;
        int s[n+1];
        s[0] = 0;
        unordered_map<int, long long> mp;
        for(int i = 0; i < n; i++) s[i+1] = s[i] + nums[i];
        for(int i = 0; i <= n; i++){
            while(s[i] < 0) s[i] = (s[i] + k) % k;
            s[i] = (s[i] + k) % k;
            mp[s[i]] ++;
        } 
        for(auto [k, v] : mp){
            // cout << k << " " << v << endl;
            cnt += (v - 1) * v / 2;
        }
        return cnt;
    }
};

LeetCode 523. 连续的子数组和

题目 示例

输入：nums = [23,2,4,6,7], k = 6
输出：true
解释：[2,4] 是一个大小为 2 的子数组，并且和为 6 。

思路

思路同上一题974,
如果map中存在 %k 后相同的历史前缀和，说明 [l,r] 区间和为k的倍数，只需再判断区间长度是否>=2即可

代码

class Solution {
public:
    bool checkSubarraySum(vector<int>& nums, int k) {
        // (s[r+1] - s[l]) % k == 0
        // s[r+1] % k == s[l] % k  &&  r + 1 - l >= 2
        int n = nums.size();
        int s[n+1];
        s[0] = 0;
        for(int i = 0; i < n; i++) s[i+1] = s[i] + nums[i];
        for(int i = 0; i <= n; i++) s[i] = s[i] % k;
        unordered_map<int, int> mp;
        for(int r = 0; r <= n; r ++){
            if(mp.count(s[r]) > 0 && r - mp[s[r]] >= 2){
                return true;
            }
            if(mp.count(s[r]) == 0) mp[s[r]] = r; // 记录pos
        }
        return false;
    }
};