yukicoder No.649 ここでちょっとQK！

「3月中に毎日1問解いて黄色になろう！」キャンペーン3日目のはずだった問題。昨晩布団の上で横になったら朝になっていたり、少し躓いたので時間がかかりました。

問題概要

2種類のクエリがQ(≦200000)個与えられるので、それを処理しましょう。 * v(≦10¹⁸)を配列に追加する * 配列の中にある要素の中でK(≦200000で、全てのクエリで共通)番目に小さな値を出力し、それを削除する。無ければ-1を出力する。

No.649 ここでちょっとQK！ - yukicoder

解説

まず普通に配列でやるとinsertにO(n)かかるので自明に死にます。ってことでmultisetを使えば1番目のクエリはクリア。

問題は2番目で、最初何も考えずに出力する値を「upper_boundの位置-最初の位置」で2分探索して、出力・削除するだけかと思ってました。実際投げてみるとTLEして、調べてみると、位置の差を取るdistance関数がrandom_iterator以外ではO(n)かかるそうで。ここで詰まりました。

バイトからの帰り道、「じゃあK番目以下かK番目以上かでデータ構造分ければよくね？」と天啓がありました。小さい順にK個は常に1つ目のsetに入れるようにして置き、残りはもう1つのsetに突っ込んでおきます。そうすると、2番目のクエリで出力するのは1つ目のsetの最後の要素です。無事間に合いますね。

https://yukicoder.me/submissions/241849

int main() {
    int q, k;  scanf("%d %d", &q, &k);
    multiset<ll> s1;
    multiset<ll> s2;
    while (q--) {
        int query; scanf("%d", &query);
        if (query == 1) {
            ll v;   scanf("%lld", &v);
            if (s1.size() < k) s1.insert(v);
            else if (v < *s1.rbegin()) {
                ll to2 = *s1.rbegin();
                s1.erase(s1.lower_bound(to2));
                s2.insert(to2);
                s1.insert(v);
            }
            else s2.insert(v);
        }
        else {
            if (s1.size() >= k) {
                ll out = *s1.rbegin();
                printf("%lld\n", out);
                s1.erase(s1.lower_bound(out));
                if (s2.size() > 0) {
                    s1.insert(*s2.begin());
                    s2.erase(s2.begin());
                }
            }
            else {
                printf("-1\n");
            }
        }
    }
    return 0;
}

感想

完全に知識不足でした。知見を得たので満足です。