LeetCode - 560. Subarray Sum Equals K

Medium

560. Subarray Sum Equals K

題目

Given an array of integers nums and an integer k, return the total number of subarrays whose sum equals to k.

A subarray is a contiguous non-empty sequence of elements within an array.

Example 1

Input: nums = [1,1,1], k = 2
Output: 2

Example 2

Input: nums = [1,2,3], k = 3
Output: 2

Constraints

• 1 <= nums.length <= 2 * 10^4
• -1000 <= nums[i] <= 1000
• -10^7 <= k <= 10^7

解題思路

這題是經典 Prefix Sum 題
目標是找有多少個 subarray 的總和等於 k

• Prefix Sum 定義：

  • sum：目前從開頭累加到現在的前綴和

• 如果某一段 subarray 的總和是 k

  • 假設目前前綴和是 sum
  • 那前面需要出現過一個 prefix sum 是 sum - k

• 因為：

  • 目前 prefix sum - 以前某個 prefix sum = k
  • 所以：
    以前某個 prefix sum = sum - k

• 使用 unordered_map<int, int> mp

  • mp[prefixSum]：代表這個 prefix sum 出現過幾次

• 初始化：

  • mp[0] = 1
  • 代表在還沒開始累加之前，prefix sum 為 0 出現過一次
  • 這樣如果從 index 0 開始的 subarray 剛好總和為 k，也可以被正確計算到

• 對每個數字 x

  • 先更新目前前綴和：
    sum += x

  • 檢查前面是否出現過 sum - k

    • 如果出現過，代表有這麼多個 subarray 可以組成總和 k
    • 所以：
      ans += mp[sum - k]

  • 最後再把目前的 sum 記錄進 map：
    mp[sum]++

• 為什麼不會重複計算？

  • 因為每次迴圈都固定把目前位置當成 subarray 的右端點
  • mp[sum - k] 只統計目前位置之前出現過的 prefix sum
  • 所以每一段 subarray 都只會在自己的右端點被計算一次

Solution Code

Time Complexity
O(n)

Space Complexity
O(n)

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        unordered_map<int, int> mp;
        mp[0] = 1;

        int sum = 0;
        int ans = 0;
        for(int x: nums){
            sum += x;
            if(mp.count(sum - k)) ans += mp[sum - k];
            mp[sum]++;
        }
        return ans;
    }
};