560. Subarray Sum Equals K

題目連結: Subarray Sum Equals K

給定一個整數陣列 nums 和一個目標值 k，請找出陣列中和為 k 的連續子陣列的個數。

範例：

輸入: nums = [1,1,1], k = 2
輸出: 2
解釋: 此題有兩個和為 2 的子陣列：[1,1] 與 [1,1]

解題思路

此題要求找出所有連續子陣列，其元素總和等於給定的目標值 k。我們需要一個高效的方法來解決此問題，時間複雜度要求為 O(n)。

方法一：暴力解法（不可取）

最直觀的想法是使用雙重迴圈，枚舉所有可能的子陣列，計算其和是否等於 k。但這種方法的時間複雜度為 O(n^2)，不符合要求。

方法二：前綴和與雜湊表

為了將時間複雜度降低到 O(n)，我們可以使用「前綴和」結合「雜湊表」的方法。

主要概念

• 前綴和（Prefix Sum）：累積計算陣列中元素的總和，sum[i] 表示從起點到索引 i 的總和。
• 雜湊表（Hash Table）：用於記錄前綴和出現的次數，方便快速查詢。

解題步驟

1. 初始化：
   • 建立一個雜湊表 prefix，鍵為前綴和，值為該前綴和出現的次數。
   • 將 prefix[0] 設為 1，表示初始前綴和為 0 出現一次。
2. 遍歷陣列：
   • 使用變數 sum 來記錄當前的前綴和，初始值為 0。
   • 初始化計數器 ans 為 0，用於統計總和為 k 的子陣列數量。
   • 對於每個元素 nums[i]：
     • 更新前綴和：sum += nums[i]。
     • 檢查 sum - k 是否在雜湊表中：
       • 如果存在，將 prefix[sum - k] 加到 ans，表示找到相應數量的子陣列總和為 k。
     • 更新雜湊表：prefix[sum] += 1，記錄當前前綴和的出現次數。
3. 返回結果：
   • 遍歷完成後，返回計數器 ans。

示例演算

假設 nums = [1, 2, 3], k = 3：

• 初始化：
  • prefix = {0:1}
  • sum = 0
  • ans = 0
• 第一輪迴圈（num = 1）：
  • sum = 0 + 1 = 1
  • sum - k = 1 - 3 = -2，prefix 不包含 -2，ans 不變。
  • 更新 prefix：prefix = {0:1, 1:1}
• 第二輪迴圈（num = 2）：
  • sum = 1 + 2 = 3
  • sum - k = 3 - 3 = 0，prefix 包含 0，ans += prefix[0] = 1，ans = 1。
  • 更新 prefix：prefix = {0:1, 1:1, 3:1}
• 第三輪迴圈（num = 3）：
  • sum = 3 + 3 = 6
  • sum - k = 6 - 3 = 3，prefix 包含 3，ans += prefix[3] = 1，ans = 2。
  • 更新 prefix：prefix = {0:1, 1:1, 3:1, 6:1}
• 結果：ans = 2

程式碼

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        const int n = nums.size();

        int ans = nums[0] == k;
        for (int i = 1; i < n; i += 1) {
            nums[i] += nums[i - 1];
            ans += (nums[i] == k);
        }

        unordered_map<int, int> prefix;
        for (int i = 0; i < n; i += 1) {
            if (prefix.count(nums[i] - k)) {
                ans += prefix[nums[i] - k];
            }

            prefix[nums[i]] += 1;
        }
        return ans;
    }
};

class Solution {
    public int subarraySum(int[] nums, int k) {
        int ans = nums[0] == k ? 1 : 0;
        for (int i = 1; i < nums.length; i += 1) {
            nums[i] += nums[i - 1];
            if (nums[i] == k) {
                ans += 1;
            }
        }

        HashMap<Integer, Integer> prefix = new HashMap<>();
        for (int i = 0; i < nums.length; i += 1) {
            if (prefix.containsKey(nums[i] - k)) {
                ans += prefix.get(nums[i] - k);
            }
            prefix.put(nums[i], prefix.getOrDefault(nums[i], 0) + 1);
        }
        return ans;
    }
}

時間與空間複雜度分析

時間複雜度：O(n)

• 遍歷陣列：需要遍歷一次長度為 n 的陣列，時間為 O(n)。
• 雜湊表操作：查詢和更新雜湊表的時間平均為 O(1)。
• 總計：O(n)。

空間複雜度：O(n)

• 雜湊表：在最壞情況下，前綴和的數量與陣列長度 n 相同，需要存儲 n 個鍵值對。
• 總計：O(n)。