【哈希表】LeetCode 560. 和爲K的子數組【中等】

給你一個整數數組 nums 和一個整數 k ，請你統計並返回 該數組中和爲 k 的子數組的個數 。

 

示例 1：

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

示例 2：

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

提示：

1 <= nums.length <= 2 * 104
-1000 <= nums[i] <= 1000
-107 <= k <= 107

【分析】

方法一：暴力解法

該題目是要找到元素和等於k的所有連續子數組的個數。那麼怎麼求一個數組的連續子數組呢。可以使用兩層循環來實現。

n = len(nums)

subnums = []
for i in range(n):
    for j in range(i, n):
        subnums.append(nums[i:j+1])

既然這樣，稍作修改，就可以求元素和爲k的連續子數組的個數了。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int: 
        n = len(nums) # 求原始數組長度
        count = 0 # 初始化子數組個數

        for i in range(n):
            for j in range(i, n):
                if sum(nums[i: j+1]) == k:
                    count += 1        

        return count

提交代碼，提示超出時間限制。

很明顯，這麼做的時間複雜度是O(n³)。因爲除了兩層for循環，還有一層求和。

其實，可以對上面方法做一些優化。因爲我們求 nums[i: j + 1] 的下一次nums[i : j + 1 + 1]，其實只比前一次多了一個元素nums[j + 1]。那麼我們就沒必要每次都從i開始算了。但是這麼做時間複雜度依然有O(n²)，依然超時。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int: 
        n = len(nums) # 求原始數組長度
        count = 0 # 初始化子數組個數

        for i in range(n):
            sum = 0
            for j in range(i, n):
                sum += nums[j]
                if sum == k:
                    count += 1
        
        return count

前綴和

什麼是前綴和：前綴和是指一個數組的某下標之前的所有數組元素的和（包含其自身）。

通常，會在前綴和首位放一個0，比如數組[1,2,3]，其前綴和是[0,1,3,6]。

前綴和通常可以幫助我們快速計算某個區間內的和。比如我們要計算i, j之間的和，那麼就是nums[i] + nums[i+1] + ... + nums[j]，可以看作是:

nums[0] + nums[1] + ... + nums[i-1] + nums[i] + nums[i+1] + ... + nums[j] - 
nums[0] + nums[1] + ... + nums[i-1]

這個式子也是preSum[j] - preSum[i-1]。

那麼，我們先遍歷一次數組，求出前綴和數組。之後這個數組可以代替我們最開始暴力解法的sum函數。但是很遺憾，這種做法複雜度依然還是O(n²)。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 要求的連續子數組
        count = 0
        n = len(nums)

        preSum = [0]

        # 求前綴和數組
        tmp = 0
        for i in range(n):
            tmp += nums[i]
            preSum.append(tmp)
        
        # 求和爲k的連續子數組，求i到j之間的和
        for i in range(1, n+1):
            for j in range(i, n+1):
                if preSum[j] - preSum[i-1] == k:  # preSum[j] - preSum[i-1]代表着在nums數組中，前j個數之和減去前i-1個數之和
                    count += 1
        
        return count

進一步優化的話，我們可以邊計算前綴和，邊統計。遍歷過程中，我們統計歷史中每一個前綴和出現的次數，然後計算到位置i(含i)的前綴和preSum減去目標k在歷史上出現過幾次，假如出現過m次，代表第i位以前不含（不含i）有m個連續子數組的和爲presum - k，這m個和爲presum - k的連續子數組，每一個都可以和presum組合成爲presum - （presum - k）= k

上述示意圖，紅色的是當前遍歷到的前綴和presum，加入它之前有兩個前綴和等於presume - k(藍色區域範圍)，那麼很明顯就會有兩個連續子數組的元素和爲k，對應橙色區域範圍。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 要求的連續子數組
        count = 0
        n = len(nums)
        preSums = collections.defaultdict(int)
        preSums[0] = 1

        presum = 0
        for i in range(n):
            presum += nums[i]

            # if preSums[presum - k] != 0
            count += preSums[presum - k] # 利用defaultdict的特性，當presum-k不存在時，返回的是0。這樣避免了判斷
            preSums[presum] += 1
        
        return count