Codeforces Round #703 (Div. 2) D - Max Median

codeforces.com

問題

You are a given an array 𝑎 of length 𝑛. Find a subarray 𝑎[𝑙..𝑟] with length at least 𝑘 with the largest median. Then, output the maximum median you can get.

解説

まず「中央値がxにできるか」という判定問題をO(N)で解くことができれば、決めうち二分探索によってこの問題をO(NlogN)で解くことができます。

ここで、判定する条件を緩和し、「中央値をx以上にできるか」とします。数列aの値をx以上である場合、1, そうでない場合-1とすれば、数列の中央値がx以上になるかは、その数列の総和で判定することができるとわかります。具体的には、総和が0より大きいなら、中央値がx以上になります。

判定条件の緩和によって、数列が決まったとき「中央値をx以上にできるか」という判定問題をO(N)で解くことができました。一方で、今回の問題では長さK以上の場合を聞かれているので少し困ります。

実は今回は、累積和による連続する数列の値の和を求める計算を考えれば、取る数列の右端を決めたとき、その右端からK以上離れた累積和の最小値が和を最大化するので、その情報を前から更新しながら持てばよいとわかります。

以上のような方針で、この問題をO(NlogN)で解くことができました。

中央値の判定問題は、以上のような「値の二値化」+「二分探索における判定条件の緩和」が絡む場合があるようです。割と上位陣が高速に説いていたので典型ということなのでしょう。

int main() {
    int N, K;
    cin >> N >> K;
    vector<int> a(N);
    rep(i,N)cin >> a[i];

    auto check = [&](int x) -> bool{
        vector<int> v(N);
        rep(i,N){
            v[i] = (a[i] >= x) ? 1: -1;
        }
        vector<int> c(N+1, 0);
        rep(i,N){
            c[i+1] = c[i] + v[i];
        }
        int minn = 0;
        for(int i = K; i <= N; i++){
            chmin(minn, c[i-K]);
            if(c[i]-minn > 0)return true;
        }
        return false;
    };

    int ok = 0, ng = *max_element(all(a)) + 1;
    while(ng - ok > 1){
        int mid = (ok + ng) / 2;
        (check(mid) ? ok: ng) = mid;
    }
    cout << ok << endl;
}