Luogu 🟣 P4062 [Code+#1] Yazid 的新生舞会

題目的難度顏色使用 Luogu 上的分級，由簡單到困難分別為 🔴🟠🟡🟢🔵🟣⚫。

🔗 🟣 P4062 [Code+#1] Yazid 的新生舞会

Problem Statement

題目簡述

給定一個長度為 $n$ 的非負整數序列 $A$。
若某個子區間的眾數出現次數嚴格大於該子區間長度的一半，則稱這個子區間是「新生舞會的」。請計算所有符合條件的子區間數量。

Constraints

約束條件

• $0 \le A_i \le n - 1$
• $type$ 只用於區分子任務，正式解法不需要依賴它

---

思路：根號分治 + 固定眾數計數

前置問題

本題是 LeetCode 🔴 3739. Count Subarrays With Majority Element II 的延伸。3739 要求計算「某個指定元素作為嚴格過半眾數」的子區間數量，而本題不指定目標值，問的是「存在任意元素成為嚴格過半眾數」的子區間總數。

注意到如果一個子區間合法，其中嚴格過半的元素必定唯一。因此每個合法區間都可以歸到唯一的「主要元素」名下，不會重複計數。因此可以往枚舉「主要元素」的方向思考，並套用 3739 的計數方法。

枚舉主要元素的問題

但直接把每個出現過的元素都當成 3739 的目標值各掃一遍，最壞會退化到 $\mathcal{O}(n^2)$，顯然是不行的，為此需要作出一些取捨。

這裡的取捨是「哪些元素值得做完整掃描，哪些元素有更有效率的方式處理」。這可以聯想到「根號分治」的思路：

1. 對於出現次數很少的元素，直接枚舉短區間即可；
2. 對於出現次數很多的元素，則可以沿用 3739 的 $\mathcal{O}(n)$ 掃描。

接下來的關鍵是：選定一個閾值 $B$，按元素在整個序列中的出現次數分成兩類，各自用最適合的方式計數。

出現次數大的元素：沿用 3739. Count Subarrays With Majority Element II 的 $O(n)$ 作法

先考慮已經指定一個元素，如何計算它作為嚴格過半元素的子區間數量，即 3739 的題目內容。這裡可以直接沿用 3739 的做法，將該元素視為 $+1$，其餘元素視為 $-1$，計算前綴和，並統計有多少對 $(i,j)$ 滿足 $i<j$ 且 $s_i<s_j$。

又由於每次加入一個新元素後，前綴和只會變化 $+1$ 或 $-1$，因此可以在掃描過程中，直接維護「歷史前綴和中有多少個小於目前前綴和」的數量。當前綴和增加或減少時，只要補上這一步新增或移出的那一層即可。

為什麼只對出現次數 $>B$ 的元素這麼做？

因為這樣的元素最多只有 $\dfrac{n}{B}$ 個，所以總時間是：

$$\mathcal{O}\left(n \cdot \frac{n}{B}\right)=\mathcal{O}\left(\frac{n^2}{B}\right).$$

出現次數小的元素：只需要看短區間

接著看全域出現次數 $\le B$ 的元素。如果它要成為某個長度為 $L$ 的子區間的嚴格過半元素，設它在這個子區間內出現了 $c$ 次，則有

$$2c > L.$$

又因為它在整個序列中最多出現 $B$ 次，所以 $c \le B$。因此合法區間長度必須滿足

$$L < 2c \le 2B.$$

也就是說，輕元素不可能成為長度 $\ge 2B$ 的子區間的嚴格過半元素。於是可以直接枚舉每個左端點，往右最多延伸 $2B-1$ 個位置，維護當前短區間中出現次數最多的元素。如果目前最大出現次數嚴格超過區間長度一半，且這個元素全域出現次數不超過 $B$，就把答案加一。

這部分每個左端點只枚舉 $\mathcal{O}(B)$ 個右端點，因此總時間是 $\mathcal{O}(nB)$。

為什麼這裡只看出現次數最多的元素就夠？

若某個區間存在嚴格過半元素，它一定也是區間內出現次數最多的元素，且嚴格過半元素唯一。因此維護目前最大頻率即可判斷這個短區間是否由輕元素貢獻。

清空計數器的細節

短區間枚舉時需要反覆維護出現次數。若每換一個左端點就重新建立或清空長度為 $n$ 的陣列，清空本身就會造成 $\mathcal{O}(n^2)$ 的額外成本，根號分治的優勢會直接消失。

本題的瓶頸不只在枚舉，也在「如何清空」。unordered_map 雖然能只記錄出現過的值，但常數較大，在本題中還是會 TLE。更好的方式是共用同一個計數陣列，另外記錄本輪被修改過的值，枚舉結束後只清掉這些位置。

這也是程式中 touched 陣列的作用：每次短區間枚舉只會碰到 $\mathcal{O}(B)$ 個位置，所以清空成本仍是 $\mathcal{O}(B)$。

複雜度平衡

兩部分合起來是

$$\mathcal{O}\left(\frac{n^2}{B}+nB\right) =\mathcal{O}\left(n\left(\frac{n}{B}+B\right)\right).$$

取 $B=\Theta(\sqrt n)$，兩項平衡，得到 $\mathcal{O}(n\sqrt n)$。

核心技巧

把「不指定主要元素」拆成兩類：

1. 出現次數 $>B$ 的元素種類少，可以逐個固定後用前綴和計數；
2. 出現次數 $\le B$ 的元素只能支配短區間，可以直接枚舉短區間並維護最大頻率。

複雜度分析

• 時間複雜度：$\mathcal{O}(n\sqrt n)$
• 空間複雜度：$\mathcal{O}(n)$

---

Code

#include <bits/stdc++.h>
using namespace std;
using i64 = long long;
#define endl '\n'

void solve() {
    int n, typ;
    cin >> n >> typ;

    vector<int> nums(n), tot(n);
    for (int& x : nums) {
        cin >> x;
        ++tot[x];
    }

    i64 ans = 0;
    int B = sqrt(n) + 1;

    // 枚舉出現次數 > B 的元素
    for (int target = 0; target < n; ++target) {
        if (tot[target] <= B) {
            continue;
        }
        // 3739. Count Subarrays With Majority Element II 的 O(n) 作法
        vector<int> cnt(2 * n + 1);
        int s = n;
        cnt[s] = 1;
        i64 lt = 0;
        for (int x : nums) {
            if (x == target) {
                lt += cnt[s];
                ++s;
            } else {
                lt -= cnt[s - 1];
                --s;
            }
            ans += lt;
            ++cnt[s];
        }
    }

    // 出現次數為 B 的元素最多只能是長度為 2B-1 區間的多數元素
    // 因此可以直接枚舉長度 < 2B 的區間，並維護出現次數最多的元素
    // unordered_map<int, int> cnt;  // TLE
    vector<int> cnt(n);
    for (int l = 0; l < n; ++l) {
        vector<int> touched;  // 紀錄被修改過的元素，方便以 O(B) 清空
        int bestCnt = 0, bestVal = -1;
        for (int r = l; r < min(n, l + 2 * B - 1); ++r) {
            int x = nums[r];
            if (cnt[x] == 0) {
                touched.push_back(x);
            }
            ++cnt[x];

            if (cnt[x] > bestCnt) {
                bestCnt = cnt[x];
                bestVal = x;
            }

            int len = r - l + 1;
            if (bestCnt * 2 > len && tot[bestVal] <= B) {
                ++ans;
            }
        }
        // 清空被修改過的元素
        for (int x : touched) {
            cnt[x] = 0;
        }
    }

    cout << ans << endl;
    return;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    solve();
    return 0;
}