ARC101 D問題 : Median of Medians - 忘れても大丈夫

問題. Median of Medians

整数列 $a = a_1, \ldots, a_n$ が与えられる． $a$ の連続部分列 $a_{l,r} = a_l, a_{l+1}, \ldots, a_{r} \, (1 \le l \le r \le n)$ の中央値を $m_{l,r}$ とする．このとき， $m = \{m_{l,r}\ \mid 1 \le l \le r \le n\}$ の中央値を求めよ．ただし，長さ $M$ の数列の中央値とは，数列を昇順にソートしたときの小さい方から $\lceil M / 2 \rceil$ 番目の要素の値である．

制約： $1 \le n \le 10^5$ ,　 $1 \le a_i \le 10^9$

解法. 二分探索 + 反点数（バブルソートの交換回数）

解けなかったので解説 [1] を参考にしています．
$a$ の連続部分列 $a_{l,r}$ に対して，集合 $A_{l,r} = \{x \in a_{l,r} \mid \lceil \frac{r - l + 1}{2} \rceil \le |\{y \in a_{l,r} \mid x \le y\}|\}$ を考える．このとき，中央値 $m_{l, r}$ と $\max_{x \in A_{l,r}} \{x\}$ の値は等しくなる．また， $x \in A_{l,r}$ ならば $y \le x$ に対して $y \in A_{l, r}$ という単調性が成り立つので， $A_{l, r}$ の最大化は二分探索で求めることができる．
現在， $a_{l, r}$ についての中央値の性質を考えたが， $m$ についても同様のことが言える． $m$ のサイズは $n (n + 1) / 2$ なので，集合 $M = \{x \in m \mid \lceil \frac{n(n+1)}{4} \rceil \le |\{y \in m \mid x \le y\}| \}$ を考えると，求める答えは $\max_{x \in M} \{x\}$ となる．したがって， $M$ に含まれるもので最大のものを二分探索で求めることを考える．
$x$ を固定して $x$ が $M$ に含まれるかどうかを判定する．連続部分列 $a_{l, r}$ に対して， $x \le m_{l, r}$ であることと， $\lceil \frac{r - l + 1}{2} \rceil \le |\{y \in a_{l,r} \mid x \le y\}|$ であることは同値である．これは， $a_{l, r}$ の要素で $x$ 未満の要素数 $s_{1}$ と $x$ 以上の要素数 $s_{2}$ に対して， $s_{1} \le s_{2}$ ということを意味している．各連続部分列に対して上の計算を高速に判定できるように累積和を用いる．長さ $n + 1$ の数列 $s$ を $s_0 = 0$ , $a_i < x$ ならば $s_i = -1$ ， $a_i \ge x$ ならば $s_i = 1$ と初期化をして累積和をとると， $x \le m_{l, r}$ と $s_{l - 1} \le s_{r}$ が同値となる．よって，各連続部分列に対して， $O(1)$ 時間で $x \le m_{l,r}$ が判定できるので， $O(n^2)$ 時間で $x \in M$ という条件 $\lceil \frac{n (n+1)}{4} \rceil \le |\{x \le m_{l,r} \mid 1 \le l \le r \le n\}|$ を計算できる．ただし，全体で $O(n^2 \log{A})$ 時間となり間に合わないので高速化を考える．
先程の計算でボトルネックとなっていたのは， $|\{x \le m_{l,r} \mid 1 \le l \le r \le n\}|$ を求めるところなので，この個数を言い換えると $|\{(l, r) \mid 0 \le l < r \le n \land s_l \le s_r \}|$ となる．これはバブルソートの交換回数を求める問題 [2] に等しく，特に求める交換回数は反点数と呼ばれている．反点数は Fenwick Tree (Binary Indexed Tree) [3] を用いると $O(n \log{n})$ 時間で求めることができる．
数列 $s_0, \ldots , s_n$ が与えられたとき $s$ の反転数を求める．まず， $0 \le i \le n$ に対して， $s_j \le s_i$ を満たす $0 \le j < i$ の数を求めるために Binary Indexed Tree(BIT) を用いる．BITは配列 $v = v_0, v_1, \ldots, v_N$ に対して次の2つのクエリを高速に計算する．
　1. $k$ 番目の要素に値 $w$ を加える（ $v_k += w$ ）： $O(1)$ 時間
　2. $v_0 + v_1 + \cdots + v_k$ を求める : $O(\log N)$ 時間
$0 \le j < i$ に対して， $v_{s_j} += 1$ というクエリが行われているときに， $v_0 + \cdots + v_{s_i}$ というクエリの結果が求める答えとなる． $i + 1$ 番目の要素に対して同様の計算を行う前に $v_{s_i} += 1$ というクエリを行うことによって仮定が満たされる．したがって， $0$ 番目の要素から順番に上の計算を行ったときのクエリの総和が反転数となる．ただし，BITの添字は $s$ の要素となり負の数になるので下駄を履かせる必要がある．ここで， $s$ の要素の絶対値の最大値は $n$ 以下となるので， $s$ の各要素に予め $n$ を加えてから計算を行う．反点数の定義から結果は変わらず，BITの添字の最大値は $2n$ 以下となる．

計算時間： $O(n \log{n} \log{A})$ （ $A$ は数列 $a$ の要素の最大値）

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

template <class T>
struct FenwickTree {
    const int n;
    std::vector<T> bit;

    FenwickTree(int _n = 0) : n(_n), bit(n + 1, 0) {}

    T sum(int r) const {
        T res = 0;
        for (r = r - 1; r >= 0; r = (r & (r + 1)) - 1) res += bit[r];
        return res;
    }
    T sum(int l, int r) const { return sum(r) - sum(l); }
    void add(int idx, const T &w) { for (; idx < n; idx |= idx + 1) bit[idx] += w; }
    int size() const { return n; }
};

class ProblemD {
public:
    ProblemD(ll _n) : n(_n), a(_n + 1), s(_n + 1, 0) {}
    void Input() { for (int i = 1; i <= n; ++i) cin >> a[i]; }

    int Solve() {
        const ll lim = ceil(n * (n + 1) / 4.0);

        ll lb = 0, ub = *max_element(a.begin(), a.end()) + 1;
        while (ub - lb > 1) {
            ll mid = 0.5 * (lb + ub), res = Check(mid);

            if (lim <= res) lb = mid;
            else ub = mid;
        }

        return lb;
    }

private:
    ll n;
    vector<int> a;
    vector<ll> s;

    ll Check(const int x) {
        for (int i = 1; i <= n; ++i) s[i] = (x <= a[i] ? 1 : -1) + s[i - 1];

        ll res = 0;
        FenwickTree<int> tree(2 * n + 1);

        for (int i = 0; i <= n; ++i) {
            res += tree.sum(0, s[i] + n + 1);
            tree.add(s[i] + n, 1);
        }

        return res;
    }
};

int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    int n;
    cin >> n;
    ProblemD p(n);
    p.Input();
    cout << p.Solve() << endl;

    return 0;
}