D - Sequence Query

[hxrxchang]

https://atcoder.jp/contests/abc241/tasks/abc241_d

問題概要

初期値が空の数列Aがある。 以下の3種類のクエリがQ個与えられるのでそれぞれ処理をせよ。

• 1 x: Aにxを追加
• 2 x k: Aのx以下の要素のうち、大きい方からk番目の値を出力する。x以下の要素がk個以上なければ -1 を出力する。
• 3 x k: Aのx以上の要素のうち、小さい方からk番目の値を出力する。x以上の要素がk個以上なければ -1 を出力する。

解き方

最初heapを使えばいいと思ったが、Aのx以下、x以上の中から何番目というのはheapではO(X)以外で取得できない。 そこでmultisetを使うといい。sortedsetの同じ値を複数入れられるバージョンで、要素の追加と並び替えを効率的に行える。 Aのx以上、x以下の値は二分探索で求め、kは最大5なので二分探索で求めたあとにk回繰り返して求めたい値を取得できる。 なおmultisetやsortedsetの内部構造は赤黒木というものが使われているらしいがよく分かっていない...

func solve() {
    q := getInt()
    ms := newMultiset[int]()

    for i := 0; i < q; i++ {
        in := getInts()
        if in[0] == 1 {
            ms.Insert(in[1])
        } else if in[0] == 2 {
            x, k := in[1], in[2]
            boundary := ms.UpperBound(x)
            // ライブラリ都合でms.Last()とms.UpperBound().Prev()で取得できる型がそれぞれ異なるためくどい書き方をしている。
            if !boundary.IsValid() {
                a := ms.Last()
                for i := 0; i < k - 1; i++ {
                    if a.IsValid() {
                        a.Prev()
                    } else {
                        break
                    }
                }
                if a.IsValid() {
                    fmt.Println(a.Value())
                } else {
                    fmt.Println(-1)
                }
            } else {
                a := boundary.Prev()
                for i := 0; i < k - 1; i++ {
                    if a.IsValid() {
                        a.Prev()
                    } else {
                        break
                    }
                }
                if a.IsValid() {
                    fmt.Println(a.Value())
                } else {
                    fmt.Println(-1)
                }
            }
        } else {
            x, k := in[1], in[2]
            a := ms.LowerBound(x).Clone()
            for i := 0; i < k - 1; i++ {
                if a.IsValid() {
                    a = a.Next()
                } else {
                    break
                }
            }
            if a.IsValid() {
                fmt.Println(a.Value())
            } else {
                fmt.Println(-1)
            }
        }
    }
}

https://atcoder.jp/contests/abc241/submissions/55038269