前綴和應用: 指定目標值的子陣列數目 Binary Subarrays With Sum_Leetcode #930

Markus Spiske on Unsplash

題目敘述

題目會給定一個陣列nums和一個目標值goal。

測試範例

Example 1:

Input: nums = [1,0,1,0,1], goal = 2
Output: 4
Explanation: The 4 subarrays are bolded and underlined below:
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]

Example 2:

Input: nums = [0,0,0,0,0], goal = 0
Output: 15

約束條件

Constraints:

• 1 <= nums.length <= 3 * 10^4

輸入陣列長度介於1~三萬。

• nums[i] is either 0 or 1.

陣列元素要嘛為0，要嘛為1。

• 0 <= goal <= nums.length

目標值goal介於0~陣列長度之間。

演算法 前綴和 + 字典

這題會用到以前學過的前綴和技巧。(和前面那題用到的觀念是類似的)

從左到右滾動依序計算前綴和，並且把每個前綴和出現的次數存到字典裡面。

如果當下在某個位置前綴和為S，而且在字典裡面又發現前面曾經看過前綴和S-goal，

那麼代表從前面到現在這個位置，存在有滿足條件的子陣列，而且子陣列的總和為goal。

例如:

給定nums=[1, 0, 1, 1], goal = 2

前綴和 = [1, 0, 2, 3]

可以找出子陣列goal = 2的部分為

[1, 0, 1],

[0, 1, 1]

[1, 1]

總共有三個子陣列滿足條件，子陣列區間總和為2。

程式碼 前綴和 + 字典

class Solution:
    def numSubarraysWithSum(self, nums, goal):
        

        ## Dictionary
        # key: prefixSum
        # value: occurrence of specific prefixSum        
        OccOfprefixSum = defaultdict(int)

        # Prefix 0 is satisfied by taking nothing
        OccOfprefixSum[0] = 1

        count, prefixSum = 0, 0

        for number in nums:

            # update prefixSum so far
            prefixSum += number

            # If we've seen "prefixSum-goal" before, then subarray with goal must exist
            count += OccOfprefixSum[prefixSum-goal]

            # Update occurrence of current prefixSum
            OccOfprefixSum[prefixSum] += 1
        
        # Total count of subarray with summation = goal
        return count

複雜度分析 前綴和 + 字典

時間複雜度:

線性掃描，所需時間為O(n)。

空間複雜度:

需要建立一個前綴和字典，所需空間為O(n)。

關鍵知識點

如果當下在某個位置前綴和為S，而且在字典裡面又發現前面曾經看過前綴和S-goal，

那麼代表從前面到現在這個位置，存在有滿足條件的子陣列，而且子陣列的總和為goal。

Reference:

[1] Binary Subarrays With Sum - LeetCode