給定一個只包含 0 和 1 的二進位陣列 nums 和整數 goal,計算有多少個連續子陣列的和恰好等於 goal(即子陣列中 1 的個數恰好為 goal)。
範例:nums = [1,0,1,0,1], goal = 2
定義前綴和 prefix[i] = nums[0] + ... + nums[i-1](prefix[0] = 0)。
若子陣列 [i+1, j] 的和等於 goal,則:
prefix[j+1] - prefix[i] == goal
⟺ prefix[i] == prefix[j+1] - goal
對每個 j,若 HashMap 中存有 prefix[j+1] - goal 的出現次數,則加入答案。
用 count[sum] 記錄每個前綴和出現的次數(初始 count[0] = 1),遍歷時先查詢再更新。
定義 atMost(k):和 ≤ k 的子陣列個數。
exactly(goal) = atMost(goal) - atMost(goal - 1)
atMost(k) 可用滑動窗口 O(n) 計算:維護左右指標,當窗口和超過 k 時收縮左邊界;右指標每移動一步,以右指標結尾的合法子陣列有 r - l + 1 個。
特別注意:goal = 0 時,atMost(-1) 需要回傳 0(而非負數),需在 atMost 中加邊界判斷。
兩種方案時間複雜度都是 O(n),HashMap 方案空間 O(n),滑窗方案空間 O(1)。
O(n),其中 n 為陣列長度。
atMost 函式被呼叫兩次,每次左右指標各只移動 O(n) 步,共 O(n)。Approach 1 (Prefix Sum + HashMap):
1. Initialize count = {0: 1}, sum = 0, result = 0
2. For each x in nums:
a. sum += x
b. result += count[sum - goal] (0 if not exists)
c. count[sum] += 1
3. Return result
Approach 2 (atMost trick):
atMost(nums, k):
1. If k < 0: return 0
2. left = 0, sum = 0, result = 0
3. For right from 0 to n-1:
a. sum += nums[right]
b. While sum > k: sum -= nums[left], left++
c. result += right - left + 1
4. Return result
numSubarraysWithSum(nums, goal):
1. Return atMost(nums, goal) - atMost(nums, goal - 1)枚舉所有子陣列 [i, j],維護滾動和並計數。
預先建立前綴和陣列,對每個 j,枚舉所有 i < j 查找 prefix[j] - prefix[i] == goal。
由於 nums 只含 0 和 1,前綴和範圍為 [0, n],可用大小 n+1 的陣列替代 HashMap。
一般整數陣列(含負數):若 nums 包含負數,滑動窗口方案無法使用(窗口和不單調),但前綴和 + HashMap 方案仍然有效,且時間仍 O(n)。這是前綴和方案相對於滑窗方案的主要優勢。
和在 [lo, hi] 範圍內的子陣列數:利用 atMost(hi) - atMost(lo - 1) 框架可直接推廣。這與 LC 795(Number of Subarrays with Bounded Maximum)的 count(right) - count(left-1) 思路完全類比,是「差分計數」的通用技巧。
和恰好為 goal 的最短/最長子陣列:若不計個數而是找最短或最長的子陣列,前綴和 + HashMap 方案只需稍作修改:記錄最早/最晚出現的前綴和索引,計算對應的子陣列長度。對二進位陣列而言,最短和為 goal 的子陣列問題等同於找連續 goal 個 1 的情況。