🔗 🟡 1248. Count Number of Nice Subarrays 1624

tags: Weekly Contest 161 前綴和(Prefix Sum) 二分搜尋(Binary Search) 雜湊表(Hash Table) 滑動窗口(Sliding Window) 數學(Math)

題意

給定一個整數陣列 numnum 和一個整數 kk。

如果一個連續子陣列上有 kk 個奇數，則該子陣列被稱為 好棒棒子陣列(nice subarray)。

返回 numnum 中 好棒棒子陣列(nice subarray) 的數量。

思路

由於我們只關注奇數的個數，因此可以將奇數的位置設為 11，偶數的位置設為 00，這樣我們就可以將問題轉換為計算子陣列中 11 的個數。

也就是說，問題可以轉換成：求出子陣列和為 kk 的子陣列數量

對於 區間求和 問題，可以使用 前綴和(Prefix Sum) 來解決。令 pre[i+1]=nums[0]+nums[1]+⋯+nums[i]=∑j=0inum[j]pre[i+1] = nums[0] + nums[1] + \cdots + nums[i] = \displaystyle\sum_{j=0}^{i} num[j]，則 [i,j][i, j] 的和為 pre[j+1]−pre[i]pre[j+1] - pre[i]。

故對於每個位置 ii ，其對答案的貢獻即為滿足 pre[j+1]−pre[i]=kpre[j+1] - pre[i] = k 的 jj 個數，即 pre[j+1]=pre[i]+kpre[j+1] = pre[i] + k 的 jj 個數。

為此可以使用 二分搜尋(Binary Search) 查找第一個以及最後一個滿足條件的 jj 位置，然後相減即可得到 jj 的個數。

  • 令 idx1=bisection_left(pre,pre[i]+k)idx1 = \text{bisection\_left}(pre, pre[i] + k) 為第一個滿足 pre[j+1]=pre[i]+kpre[j+1] = pre[i] + k 的位置。
  • 令 idx2=bisection_right(pre,pre[i]+k)−1idx2 = \text{bisection\_right}(pre, pre[i] + k) - 1 為最後一個滿足條件的位置，即第一個滿足 pre[j+1]>pre[i]+kpre[j+1] > pre[i] + k 的位置的前一個位置。
  • jj 的個數即為 idx2−idx1+1idx2 - idx1 + 1，累加到答案中。

複雜度分析

  • 時間複雜度：O(nlog⁡n)\mathcal{O}(n \log n)，需要遍歷 nn 個元素，每次二分搜尋需要 log⁡n\log n 時間。
  • 空間複雜度：O(n)\mathcal{O}(n)，前綴和陣列需要的空間。
1
2
3
4
5
6
7
8
9
10
class Solution:
    def numberOfSubarrays(self, nums: List[int], k: int) -> int:
        n = len(nums)
        s = [0] * (n + 1)
        for i, x in enumerate(nums):
            s[i + 1] = s[i] + x % 2
        ans = 0
        for i in range(n):
            ans += bisect_right(s, s[i] + k) - bisect_left(s, s[i] + k)
        return ans
1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution {
public:
    int numberOfSubarrays(vector<int>& nums, int k) {
        int n = nums.size();
        vector<int> s(n + 1);
        for (int i = 0; i < n; i++)
            s[i + 1] = s[i] + nums[i] % 2;
        int ans = 0;
        for (int i = 0; i < n; i++) {
            ans += upper_bound(s.begin(), s.end(), s[i] + k) - lower_bound(s.begin(), s.end(), s[i] + k);
        }
        return ans;
    }
};

方法二：前綴和(Prefix Sum) + 雜湊表(Hash Table)

這個解法的思路與方法一類似，也是利用 前綴和(Prefix Sum) 來計算區間和。注意到在方法一中我們是考慮左端點，因此需要用二分搜尋來查找符合條件的右端點數量。但可以反過來考慮，我們可以考慮右端點，然後 紀錄符合條件的左端點數量 。

為此，我們需要使用一個 雜湊表(Hash Table) cntcnt 來記錄前綴和值的出現次數，並在當前位置找到滿足條件的左端點數量。由於出現次數最多為 nn ，因此也能使用一個長度為 n+1n+1 的陣列來記錄。

具體做法如下：

  1. 初始化前綴和 ss 和一個計數陣列 cntcnt，將 cnt[0]cnt[0] 設為 11 。
  2. 遍歷陣列 numsnums，根據 nums[i]nums[i] 的奇偶性更新前綴和 ss。
  3. 如果 s≥ks \geq k，則表示從某個位置到當前位置的子陣列和為 kk。我們需要找出從哪裡開始到當前位置的子陣列個數，這個個數即為 cnt[s−k]cnt[s-k]。
  4. 更新 cnt[s]cnt[s]，表示前綴和為 ss 的位置個數。
  5. 最後返回答案 ansans。

複雜度分析

  • 時間複雜度: O(n)\mathcal{O}(n), 只需要遍歷陣列一次。
  • 空間複雜度: O(n)\mathcal{O}(n), 使用了一個長度為 n+1n+1 的陣列。
1
2
3
4
5
6
7
8
9
10
11
class Solution:
    def numberOfSubarrays(self, nums: List[int], k: int) -> int:
        n = len(nums)
        ans = s = 0
        cnt = [1] + [0] * n
        for x in nums:
            s += x & 1
            if s >= k:
                ans += cnt[s - k]
            cnt[s] += 1
        return ans
1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution {
public:
    int numberOfSubarrays(vector<int>& nums, int k) {
        int n = nums.size(), ans = 0, s = 0;
        vector<int> cnt(n + 1);
        cnt[0] = 1;
        for (int x : nums) {
            s += x & 1;
            if (s >= k) ans += cnt[s - k];
            cnt[s]++;
        }
        return ans;
    }
};

方法三：滑動窗口(Sliding Window) + 數學(Math)

由於子陣列中必須要有 kk 個奇數，因此我們可以考慮其最左邊的奇數位置和最右邊的奇數位置，左方能往左延伸長度即為左方的偶數個數 lcntlcnt，右方能往右延伸長度即為右方的偶數個數 rcntrcnt，根據 乘法原理 ，滿足條件的子陣列個數即為 (lcnt+1)×(rcnt+1)(lcnt + 1) \times (rcnt + 1)。

為此需要紀錄每個奇數的位置，令 oddodd 為所有奇數的位置，則 odd[i]odd[i] 表示第 ii 個奇數的位置。並透過枚舉最左邊的奇數位置來得出對應的右邊奇數位置，然後計算左右偶數個數即可得出該位置對答案的貢獻。

具體步驟如下：

  1. 初始化一個陣列 oddodd 來存儲所有奇數的位置，則第 ii 個奇數的位置為 odd[i]odd[i]。
    • 為了後續計算方便，可以在 oddodd 陣列的前面加上 −1-1 和在後面加上 nn。
  2. 初始化答案 ans=0ans = 0。
  3. 遍歷 oddodd 陣列，對於每個位置 ii，計算左右偶數個數 lcntlcnt 和 rcntrcnt。
    • 最左邊的奇數位置是第 ii 個奇數，則 lcnt=odd[i]−(odd[i−1] if i>0 else −1)−1lcnt = odd[i] - (odd[i-1] \text{ if } i > 0 \text{ else } -1) - 1。
    • 最右邊的奇數位置是第 j=i+k−1j = i + k - 1 個奇數，則 rcnt=(odd[j+1] if j+1<len(odd) else n)−odd[j]−1rcnt = (odd[j+1] \text{ if } j + 1 < \text{len(odd)} \text{ else } n) - odd[j] - 1。
  4. 將 (lcnt+1)×(rcnt+1)(lcnt + 1) \times (rcnt + 1) 加到答案中。
    • 這裡加 11 是因為左右偶數個數都可以為 00，即左邊可以再添加 [0,lcnt][0, lcnt] 個偶數，右邊可以再添加 [0,rcnt][0, rcnt] 個偶數。
  5. 最後返回答案即可。

複雜度分析

  • 時間複雜度：O(n)\mathcal{O}(n)。
    • 計算 oddodd 陣列的時間複雜度是 O(n)\mathcal{O}(n)。
    • 枚舉最左邊的奇數位置的時間複雜度是 O(n)\mathcal{O}(n)。
  • 空間複雜度：O(n)\mathcal{O}(n)，使用了一個長度為 nn 的 oddodd 陣列。
1
2
3
4
5
6
7
8
9
10
11
class Solution:
    def numberOfSubarrays(self, nums: List[int], k: int) -> int:
        n = len(nums)
        odd = [i for i, x in enumerate(nums) if x & 1]
        ans = 0
        for i in range(len(odd) - k + 1): # 枚舉最左邊的奇數
            j = i + k - 1 # 最右邊的奇數
            lcnt = odd[i] - (odd[i-1] if i > 0 else -1) - 1 # 左邊的偶數個數
            rcnt = (odd[j+1] if j + 1 < len(odd) else n) - odd[j] - 1 # 右邊的偶數個數
            ans += (lcnt + 1) * (rcnt + 1)
        return ans
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Solution {
public:
    int numberOfSubarrays(vector<int>& nums, int k) {
        int n = nums.size(), ans = 0;
        vector<int> odd;
        odd.push_back(-1); // dummy
        for (int i = 0; i < n; i++)
            if (nums[i] & 1) odd.push_back(i);
        odd.push_back(n); // dummy
        for (int i = 1; i + k < odd.size(); i++) {
            int j = i + k - 1;
            ans += (odd[i] - odd[i - 1]) * (odd[j + 1] - odd[j]);
        }
        return ans;
    }
};

寫在最後

Cover photo is created by @たろたろ, thanks for their work!