leetcode-滿足連續子數組加和等於目標值的子數組個數 思路與代碼

問題描述

問題鏈接：https://leetcode.com/problems/subarray-sum-equals-k/
給定一個數組，請找到有多少個連續子數組加和等於給定目標數值？

leetcode，medium
560：Subarray Sum Equals K

問題分析

首先想到的是兩件事情：

1. 兩層for循環遍歷的解法；
2. 兩層for循環遍歷大概率會出現TLE的錯誤，哈哈，不是廢話嘛

因此，我們需要將O(n^2)的解法從時間上優化成O(n)的解法，這一般會增加空間上的開銷，即“空間開銷換時間開銷”。這裏參考的網上大神的常規做法，如下:

⚠️：下面提到的累加和，在不做特別說明的情況下，都是指從數組開頭到當前位置處的所有元素的累加和。

1. 申請字典數據結構m，初始時m中的元素爲{0: 1}（後面會解釋原因），同時初始化滿足條件的子數組的數量爲res=0；
2. 記錄以當前位置編號爲i的元素作爲結尾的累加和sum；
3. 從字典m中查詢key= sum-k 對應的值m[sum-k]，表示在位置編號爲i的元素之前已經出現了m[sum-k]個元素作爲結尾的累加和爲sum-k的情況。例如以位置編號爲i-5的元素結尾的累加和爲sum-k ，那麼從位置i-4到i之間的所有子數組的加和一定爲k，剛好滿足條件。此時以位置編號爲i的元素作爲結尾的子數組（不一定從數組開頭開始計算）加和等於k的情況一共是m[sum-k]種，此時res的更新只需要在原來的res基礎上再加m[sum-k]即可，即res = res + m[sum-k]；
4. 對於累加和sum，其在字典結構m中以sum作爲key的值應該自加1。

在步驟3中，當sum-k剛好等於0時，說明此時從數組開頭到當前位置所有元素累加和剛好等於k，滿足條件，此時res就要更新了。但是當第一次碰到這種情況時，res = res + m[sum-k]中的 m[sum-k]=m[0]=0（字典數據結構對於新的key的初始化其value爲0）導致res無法更新，所以我們人爲的在最開始時初始化{0: 1}即可。

問題解法

直接看代碼吧，結合問題分析中的內容，代碼應該很容易理解了。

def subarray_sum_equals_k(nums, target):
    l = len(nums)
    res = 0
    m = {0: 1}
    sum = 0
    for i in range(l):
        sum += nums[i]
        res += m.get(sum - k, 0)
        m[sum] = m.get(sum, 0) + 1
    return res

測試代碼如下：

if __name__ == "__main__":
    nums = [1, 1, 1]
    k = 2
    print(subarray_sum_equals_k(nums, k))