AtCoder Regular Contest 101 ARC 101 D - Median of Medians

問題

長さNの数列が与えられる。この中の全ての連続部分列における中央値の中央値を求めよ。

問題へのリンク

方針

単調性

二分探索する。 中央値がX以上となる区間の数が全ての区間のうち半分以上あるか？を確認する。 この関係には単調性がある。つまり、あるXのときは半分以上あるが、X+1では半分未満という境界があるということである。この境界を調べるために二分探索をする。

判定問題　区間の中央値がX以上になる条件

あるXが与えられたとき、以下の判定問題を考えたい。

中央値がX以上となる区間の数が全ての区間のうち半分以上あるか？

判定するために必要なのは、要素がX以上か、未満かという大小関係である。

• 要素がX以上の場合1,

• そうではない場合-1

という数列に書き換えれば、区間の和が0以上になることが条件になる。ぴったり半分なら和が0になるし、奇数個なら和は0にはならないが、1以上なら中央値はX以上である。かなり綺麗な言い換え。

区間列挙の解決策

判定問題は綺麗に言い換えることができた。しかし・・・

区間の和を全て試すわけにもいかない。どうすれば高速に区間の和を考えることができるだろうか？ ここで、区間の和は、累積和をとったとき、左端と右端を決めればその差で高速に求められることを思い出したい。

これを応用すると、右端を全て見ながら（とでもする）、左側に以下の値がいくつあるかがわかればよい。これは、prefixの和であり、BITの出番である。右端を進めながらメモしていくことでrより左側に以下の要素がいくつあるかをで求められる。

以上により、で問題を解くことができた。

実装

累積和が負の数になるので、BITの長さをくらい取って値を突っ込むときに+Nするようにするときれいに実装できる。

提出のリンク

#include <bits/stdc++.h>
using namespace std;
#define rep(i,N) for(int i=0;i<int(N);++i)
#define print(v) { cerr<<#v<<": [ "; for(auto _ : v) cerr<<_<<", "; cerr<<"]"<<endl; }
typedef long long ll;
const ll INFLL = (ll)1e18+1;

template<typename T>
class BIT{
  public:
    int N;
    vector<T> data;
    BIT(T _N):N(_N){
        data.assign(N, 0);
    };
    
    // a is 1-indexed
    void add(int a, T w){
        for(int x = a; x <= N; x += x & -x)data[x] += w;
    }
    // 1-indexed sum of prefix [0, a]
    T sum(int a){
        T res = 0;
        for(int x = a; x > 0; x -= x & -x)res += data[x];
        return res;
    }
    // 1-indexed sum of range [l, r]
    T sum(int l, int r){return sum(r) - sum(l-1);}

    // 0-indexed add
    void add0(int a, T w){add(a + 1, w);}
    // 0-indexed sum
    T sum0(int a){return sum(a + 1);}
    // 0-indexed sum of range
    T sum0(int l, int r){return sum0(r) - sum0(l-1);}
    // show the value
    void debug(){print(data);}
};

int main() {

    ll N;
    cin >> N;
    ll half = (N+1) * N / 4;
    vector<ll> a(N);
    rep(i,N)cin>>a[i];
    if(N == 1){
        cout << a[0] << endl;
        return 0;
    }
    ll ok = 0;
    ll ng = INFLL;

    // 中央値がX以上になる区間が半分以上あるか？
    auto check = [&](ll X){
        vector<ll> b(N);
        rep(i,N)b[i] = ((a[i] >= X) ? 1: -1);
        vector<ll> cum(N+1);
        rep(i,N)cum[i+1] = cum[i] + b[i];
        BIT<ll> Tree(2 * N + 10);
        ll cnt = 0;
        rep(i,N+1){
            cnt += Tree.sum(cum[i]+N+1);
            Tree.add(cum[i]+N+1, 1);
        }
        return half  <= cnt;
    };

    while(ng - ok  > 1){
        ll mid = (ok + ng) / 2;
        if(check(mid))ok = mid;
        else ng = mid;
    }

    cout << ok << endl;
}