はまやんはまやんはまやん

hamayanhamayan's blog

L番目のK番目の数 (LthKthNumber) [JOI/IOI 第17回日本情報オリンピック 予選 F]

https://atcoder.jp/contests/joi2018yo/tasks/joi2018_yo_f

解法

https://atcoder.jp/contests/joi2018yo/submissions/8137835

K番目に小さな整数を愚直に全列挙することを考えてみる。
N=2*10^5の場合は無理だが、N=4000なら可能であるため、小課題2は全列挙してL番目を出すことを考える。
最大ケースでは全列挙は無理なのでなんらかの方法でL番目を特定するが、これは典型がある。
「check(x) := K番目に小さな整数列の中でx以下の数が何個あるか」
これを使って二分探索でL番目を特定していく。
経験が無いと、まずここまでが厳しいかもしれない。
 
次の問題はcheck(x)の計算である。
この問題はO(N)か最悪O(NlogN)くらいではやりたい。
よくあるのが右端Rを全探索して、満たす左端を数え上げるというテクがあるが、こっちの方面で考えてみよう。
[?,R]の範囲でK番目に小さい整数がx以下である個数は何個かと言う問題を考える。
言い換えると「[?,R]の範囲でx以下の数がK個以上となる区間はいくつあるか」という問題になる
「f(L,R,x) := [L,R]の範囲でx以下の数は何個あるか」を用意すると「f(L,R,x) = f(0,R,x) - f(0,L-1,x)」
つまり、
K ≦ f(L,R,x)
K ≦ f(0,R,x) - f(0,L-1,x)
f(0,L-1,x) ≦ f(0,R,x) - K
となり、上記の式を満たすLの個数を数え上げればいい。
BITと累積和を使えば以上の問題は解決できる。
 
公式解説では、BITを使っている部分を尺取り法で行っている。
確かにRを単調増加させていくと、Lも単調増加していく。
BIT解法でも間に合うが尺取り法の方が計算量がいい。

int N, K; ll L; int A[201010];
//---------------------------------------------------------------------------------------------------
int cnt[201010];
ll check(int x) {
    BIT<ll> bit(N + 1);
    
    bit.add(0, 1);
    rep(i, 0, min(N,K - 1)) {
        if(i) cnt[i] = cnt[i - 1];
        else cnt[i] = 0;
        if (A[i] <= x) cnt[i]++;
    }

    ll res = 0;
    rep(R, K - 1, N) {
        cnt[R] = cnt[R - 1];
        if (A[R] <= x) cnt[R]++;
        bit.add(cnt[R - K + 1], 1);

        // f(0,L-1,x) ≦ f(0,R,x) - K
        if(0 <= cnt[R] - K + 1) res += bit.get(0, cnt[R] - K + 1);
    }
    return res;
}
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> K >> L;
    rep(i, 0, N) cin >> A[i];

    //rep(i, 2, 10) printf("%d -> %lld\n", i, check(i));

    int ng = 0, ok = N;
    while (ng + 1 != ok) {
        int x = (ng + ok) / 2;
        if (L <= check(x)) ok = x;
        else ng = x;
    }
    cout << ok << endl;
}