LeetCode 🟡 1248. Count Number of Nice Subarrays

🔗 🟡 1248. Count Number of Nice Subarrays 1624

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

題意

給定一個整數陣列 $num$ 和一個整數 $k$。

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

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

思路

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

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

方法一：前綴和(Prefix Sum) + 二分搜尋(Binary Search)

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

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

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

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

複雜度分析

• 時間複雜度：$\mathcal{O}(n \log n)$，需要遍歷 $n$ 個元素，每次二分搜尋需要 $\log n$ 時間。
• 空間複雜度：$\mathcal{O}(n)$，前綴和陣列需要的空間。

PythonC++

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

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) $cnt$ 來記錄前綴和值的出現次數，並在當前位置找到滿足條件的左端點數量。由於出現次數最多為 $n$ ，因此也能使用一個長度為 $n+1$ 的陣列來記錄。

具體做法如下：

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

複雜度分析

• 時間複雜度: $\mathcal{O}(n)$, 只需要遍歷陣列一次。
• 空間複雜度: $\mathcal{O}(n)$, 使用了一個長度為 $n+1$ 的陣列。

PythonC++

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

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)

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

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

具體步驟如下：

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

複雜度分析

• 時間複雜度：$\mathcal{O}(n)$。
  • 計算 $odd$ 陣列的時間複雜度是 $\mathcal{O}(n)$。
  • 枚舉最左邊的奇數位置的時間複雜度是 $\mathcal{O}(n)$。
• 空間複雜度：$\mathcal{O}(n)$，使用了一個長度為 $n$ 的 $odd$ 陣列。

PythonC++

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

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 generated by @たろたろ, thanks for their work!