🔗 🔴 3739. Count Subarrays With Majority Element II

Problem Statement

題目簡述

給定整數陣列 nums 和整數 target，請計算有多少個非空子陣列滿足：target 在該子陣列中出現的次數嚴格大於子陣列長度的一半，也就是 target 是這個子陣列的多數元素。

Constraints

約束條件

$1 \leq \texttt{nums.length} \leq 10^5$
$1 \leq \texttt{nums[i]} \leq 10^9$
$1 \leq \texttt{target} \leq 10^9$

思路：枚舉右維護左

前置知識

請先閱讀 LeetCode 🟡 3737. Count Subarrays With Majority Element I 的筆記，了解「將 target 轉成 +1、其他元素轉成 -1，問題變成統計前綴和遞增配對」的等價轉換。

主要結論為：將 $nums[i]$ 根據是否為 target 轉成 $+1$ 或 $-1$ 後，對其做前綴和。對於子陣列 $[l,r]$ ，如果前綴和差值滿足 $s_{r+1}-s_l>0$ ，則滿足條件。

差別在於這版的陣列長度到 $10^5$ ，不能再枚舉所有左右端點。

修改下標定義，令 $i = l,\quad j = r + 1$ ，則對於 $0 \le i < j \le n$ ，只要滿足 $s_i < s_j$ 就符合題意。也就是說，問題變成統計有多少對 $(i,j)$ 滿足 $i<j$ 且 $s_i<s_j$ 。

那麼問題來了，我們能不能在枚舉前綴和右端點

j

時，維護有多少個左端點

i

滿足

i<j

且

s_i<s_j

？

如果可以做到 $O(\log n)$ 或 $O(1)$ 維護的話，就能把整體複雜度降到 $O(n\log n)$ 甚至是 $O(n)$ 。

條件是嚴格小於，不是小於等於。若兩個前綴和相等，區間和為

0

，表示 target 和非 target 數量相同，並不符合「嚴格超過一半」。

二維偏序轉換

每個前綴和可以看成一個點 $(j, s_j)$ ，第一維是前綴下標，第二維是前綴和值。要計算的配對條件是：對目前點 $(j, s_j)$ ，找出有多少歷史點 $(i, s_i)$ 滿足 $i<j$ 且 $s_i<s_j$ 。

這就是二維偏序計數，類似求逆序對。由於我們按照 $j$ 從小到大掃描，第一維 $i<j$ 已經天然滿足，剩下只要用資料結構維護歷史前綴和值，並查詢其中有多少個值小於目前的 $s_j$ 。

方法一：有序容器上二分

最直接的做法是維護已經出現過的前綴和集合。掃描到目前的前綴下標 $j$ 時，滿足 $s[i] < s[j]$ 的歷史前綴和 $s[i]$ ，就是集合中所有小於 $s[j]$ 的元素。

有序容器支援兩件事：

查詢有多少歷史前綴和小於目前前綴和。
把目前前綴和加入歷史集合，供後面的前綴下標使用。

初始時需要先放入前綴和 $s_0=0$ ，對應空前綴。這樣當子陣列從陣列開頭開始時，也能被正常計入。

複雜度分析

時間複雜度： $\mathcal{O}(n\log n)$ ，每個位置做一次二分查詢與插入。
空間複雜度： $\mathcal{O}(n)$ ，需要儲存歷史前綴和。

方法二：二維偏序，逆序對問題

從配對條件 $i<j$ 且 $s_i<s_j$ 來看，這就是一個二維偏序計數：一維是下標先後，另一維是前綴和值大小。由於掃描順序已經保證歷史前綴和都來自目前位置左側，剩下只要快速統計「值小於目前前綴和」的歷史數量。

前綴和的範圍只會落在 $[-n,n]$ ，需要用偏移量把它映射到 $[1,2n+1]$ ，然後用樹狀陣列維護每個前綴和值出現了幾次。每次查詢小於目前值的總出現次數，再把目前值加入資料結構。

為什麼可以不用離散化？

一般逆序對或偏序計數會先離散化數值，但這裡每個位置只會貢獻 $+1$ 或 $-1$ ，前綴和值域天然被限制在 $[-n,n]$ 。因此用一個固定偏移量就能把所有可能值映射到樹狀陣列下標。

複雜度分析

時間複雜度： $\mathcal{O}(n\log n)$ ，每個位置做一次樹狀陣列查詢與更新。
空間複雜度： $\mathcal{O}(n)$ ，值域大小為 $2n+1$ 。

方法三：動態維護小於目前值的數量

注意到每次加入一個新元素後，前綴和只會變化 $+1$ 或 $-1$ 。這意味著：當我們從上一個前綴和走到下一個前綴和時，「歷史前綴和中有多少個值小於目前值」這個答案，不會整體重算，只會發生局部變化。

因此可以在掃描過程中，直接維護「歷史前綴和中有多少個小於目前前綴和」的數量。當前綴和增加或減少時，只要補上這一步新增或移出的那一層即可。

若目前元素等於 target，前綴和從 $v$ 變成 $v+1$ 。原本等於 $v$ 的歷史前綴和，現在也變成小於新值，所以小於數量要加上「歷史中等於 $v$ 的個數」。
若目前元素不等於 target，前綴和從 $v$ 變成 $v-1$ 。原本等於 $v-1$ 的歷史前綴和，不再小於新值，所以小於數量要減去「歷史中等於 $v-1$ 的個數」。

這樣就不需要有序容器或樹狀陣列了，只要用計數陣列記錄每個前綴和值出現次數，並維護目前的小於數量即可。

核心技巧

核心技巧是利用前綴和每次只變動 $1$ 的特性，把「查詢有多少歷史值小於目前值」從資料結構查詢改成增量維護。這是從 $O(n\log n)$ 降到 $O(n)$ 的關鍵。

複雜度分析

時間複雜度： $\mathcal{O}(n)$ ，每個位置只做常數次計數更新。
空間複雜度： $\mathcal{O}(n)$ ，需要記錄 $[-n,n]$ 範圍內的前綴和出現次數。

方法四：分治

要怎麼想到分治？回到配對條件 $i<j$ 且 $s_i<s_j$ 。前面的方法都是按照 $j$ 從左到右掃描，然後用資料結構回答「左邊有多少個更小的前綴和」。換個角度看，這其實是在統計一堆跨下標的配對，而這類配對計數也常常可以用 merge sort 分治處理。

具體來說，把前綴和陣列按照下標切成左右兩半。合法配對只會有三種來源：完全在左半部、完全在右半部、以及左端點在左半部且右端點在右半部。前兩種交給遞迴處理，剩下的跨區間配對，因為左半部下標一定小於右半部下標，所以只需要判斷前綴和值是否滿足 $s_i<s_j$ 。

跨區間的統計正好可以放進 merge sort 的合併過程。程式中會先讓左右兩段各自排好序，然後以「由大到小」的方式合併：

若左側目前值大於等於右側目前值，就先放左側，因為它不能和這個右側值形成合法配對。
若左側目前值小於右側目前值，代表左側從目前位置到中點的所有值都小於這個右側值，於是可以一次累加這些配對數量。

這種寫法本質上和「逆序對」很像，只是逆序對通常統計的是 $A_i>A_j$ ，這裡改成統計 $A_i<A_j$ 。因此標準 merge sort 求逆序對時，常見寫法會維持由小到大排序；而這裡為了方便一次統計左側有多少值小於右側目前值，合併方向改成由大到小。只要把合併方向與計數條件對應好，就能在分治過程中把所有合法配對算完。

可以把這個做法理解成「用 merge sort 計算前綴和陣列中的遞增對/逆序對數量」

複雜度分析

時間複雜度： $\mathcal{O}(n\log n)$ ，每一層分治都會線性合併一次，總共有 $\log n$ 層。
空間複雜度： $\mathcal{O}(n)$ ，合併時需要暫存陣列，遞迴深度為 $\mathcal{O}(\log n)$ 。

Code

方法一：有序容器上二分

class Solution:
    def countMajoritySubarrays(self, nums: List[int], target: int) -> int:
        ans = s = 0
        sl = SortedList()
        sl.add(0)
        for x in nums:
            s += 1 if x == target else -1
            ans += sl.bisect_left(s)
            sl.add(s)
        return ans

方法二：二維偏序，逆序對問題

class BIT:  # PURQ, 1-based
    __slots__ = ["tree"]

    def __init__(self, n: int):
        self.tree = [0] * (n + 1)

    def add(self, k: int, x: int) -> None:  # 令 nums[k] += x
        while k < len(self.tree):
            self.tree[k] += x
            k += k & -k

    def preSum(self, k: int) -> int:  # 求 nums[:k+1] 之和
        res = 0
        while k > 0:
            res += self.tree[k]
            k -= k & -k
        return res

    def query(self, l: int, r: int) -> int:  # 求 nums[l:r+1] 之和
        if l > r:
            return 0
        return self.preSum(r) - self.preSum(l - 1)


class Solution:
    def countMajoritySubarrays(self, nums: List[int], target: int) -> int:
        n = len(nums)

        ans = 0
        s = n + 1  # offset 為 n + 1，起始為 0 + offset = n + 1
        bit = BIT(2 * n + 1)  # [-n, n] -> [1, 2n + 1]
        bit.add(s, 1)
        for x in nums:
            s += 1 if x == target else -1
            ans += bit.query(1, s - 1)
            bit.add(s, 1)
        return ans

方法三：動態維護小於目前值的數量

class Solution:
    def countMajoritySubarrays(self, nums: List[int], target: int) -> int:
        n = len(nums)
        cnt = [0] * (2 * n + 1)  # [-n, n] => [0, 2n]

        # s = cnt_target - cnt_other
        s = n  # offset 為 n，起始為 0 + offset = n
        cnt[s] = 1
        # lt = 有多少個前綴下標 i 滿足 i < j 且 s[i] < s[j]
        ans = lt = 0
        for x in nums:
            if x == target:
                lt += cnt[s]
                s += 1
            else:
                lt -= cnt[s - 1]
                s -= 1
            ans += lt
            cnt[s] += 1
        return ans

方法四：分治

class Solution4:
    def countMajoritySubarrays(self, nums: List[int], target: int) -> int:
        n = len(nums)

        A = list(accumulate((1 if x == target else -1 for x in nums), initial=0))

        def cdq(left, right):
            if left >= right:
                return 0
            mid = (left + right) // 2
            res = cdq(left, mid) + cdq(mid + 1, right)

            # 如果左右兩側已經有序，則不需要合併，且不會產生逆序對
            if A[mid] >= A[mid + 1]:
                return res

            # 使用 Merge Sort 將左右兩側合併（從大到小排列），並計算逆序對數量
            i, j = left, mid + 1
            tmp = []
            while i <= mid or j <= right:
                if j > right or i <= mid and A[i] >= A[j]:
                    tmp.append(A[i])
                    i += 1
                else:
                    tmp.append(A[j])
                    # 逆序對數量，即滿足原始陣列中 i < j 且 A[i] < A[j] 的數量（本題為由大到小）
                    res += mid - i + 1
                    j += 1
            A[left : right + 1] = tmp
            return res

        return cdq(0, n)