No.3298 K-th Slime （tatyam set を使う解法）

問題概要

問題URL

問題文

$N$ 個のスライムが入っている箱があります。各スライムの大きさはそれぞれ $A_1, A_2, \cdots, A_N$ です。

$Q$ 個のクエリが与えられるので、順に処理してください。

• 1 x : 大きさ $x$ のスライムを新しく生成し、箱の中に入れる。
• 2 y : 箱の中の $K$ 番目に小さいスライムを $1$ つ取り出す。そのスライムの大きさを $y$ 大きくし、箱の中に戻す。
• 3 : 箱の中の $K$ 番目に小さいスライムの大きさを出力する。

制約

• $1 \leq K \leq N \leq 10^5$
• $1 \leq Q \leq 10^5$
• $1 \leq A_i \leq 10^9$
• $1 \leq x, y \leq 10^9$
• 入力はすべて整数である

解説

tatyam さんの SortedMultiSet を貼ることで、比較的楽に解くことができます。

multiset を用いて、各クエリを次のように処理すれば良いです。はじめ multiset = SortedMultiSet(A) で数列 $A$ として初期化します。

• クエリ $1$ : multiset.add で $x$ を追加する
• クエリ $2$ : 以下の手順で処理する
  • a = multiset.pop で $K$ 番目に小さい値 $a$ をポップする
  • multiset.add で $a + y$ を追加する
• クエリ $3$ : multiset[x] で $K$ 番目に小さい値を取得する

このようにすることで、計算量 $O((N + Q) \sqrt N)$ で解くことができました。

# SortedsMultiSet 省略

N, K, Q = map(int, input().split())
K -= 1
A = SortedMultiset(map(int, input().split()))

for _ in range(Q):
    cmd, *query = map(int, input().split())

    match cmd:
        case 1:
            x = query[0]
            A.add(x)
        case 2:
            y = query[0]
            A.add(A.pop(K) + y)
        case 3:
            print(A[K])