大きい方/小さい方からtopKを管理しながらなんかやるモジュールを書いた

ごり〜　（こんばんは)

なんか書きました
大したものじゃないんですけどたまによくデータ構造系のやつで「heapq2本持ちながらやる」みたいによく言われるやつです
実装の都合でheapq2本ではなくtatyamさんの「SortedMultiset」を2本持つ形で実装しています
heapqじゃないので速度は期待しないでください (代わりにSortedMultisetの部分は高機能です)
コンテスト中に書いてもいいぐらいのコードだけど、たまに出てくる気がするのでペタりできたほうが嬉しいよね

今回ご紹介するのは

こちらです、 TopKSortedMultiset
詳しくは上記リンク先に書きましたが、自由に追加/削除が O(√N) ででき、任意の時点でtopK個の和などが取れます
和じゃないものが必要になったら適宜改造して使ってください☆★☆★(なぞの星)

使い方実例、verify

ABC281 E - Least Elements

各 から までのうち小さい方から個の和を高速に求める問題。
切れ味抜群です。

Submission #53027580 - AtCoder Beginner Contest 281

N,M,K = map(int,input().split())
A = list(map(int,input().split()))
ts = TopKSortedMultiset(K, A[:M], mode_max=False)
ans = [ts.sum_topk]
for i in range(M,N):
    ts.add(A[i])
    ts.discard(A[i-M])
    ans.append(ts.sum_topk)
print(*ans)

ABC306 E - Best Performances

クエリごとに をに書き換えていき、大きい方から個の和を求める問題。
これも切れ味抜群。

Submission #53027660 - Toyota Programming Contest 2023#3（AtCoder Beginner Contest 306）

N,K,Q = map(int,input().split())
XY = [tuple(map(int,input().split())) for _ in range(Q)]
A = [0] * N
ts = TopKSortedMultiset(K, [0]*N)
for x,y in XY:
    x -= 1
    ts.discard(A[x])
    ts.add(y)
    A[x] = y
    print(ts.sum_topk)

Donutsプロコンチャレンジ2015 - ドーナツの箱詰め

このモジュールを作ろうと思ったきっかけになった問題。
やることが整理できていれば実装は楽になるはず

Submission #53003393 - Donutsプロコンチャレンジ2015

N,K = map(int,input().split())
C = list(map(int,input().split()))
Q = int(input())
qs = [C[int(input())-1] for i in range(Q)]

ss = SortedSet(C)
ini_arr = []
p = None
for c in ss:
    if p is not None:
        ini_arr.append(c-p)
    p = c
ts = TopKSortedMultiset(K-1, ini_arr)
print(ts.sum_other)

for q in qs:
    ss.discard(q)
    l = ss.lt(q)
    r = ss.gt(q)
    if l is None:
        ts.discard(r-q)
    elif r is None:
        ts.discard(q-l)
    else:
        ts.discard(r-q)
        ts.discard(q-l)
        ts.add(r-l)
    print(ts.sum_other)

できそうでTLEしたやつ (ざんねん)

今後の課題

ABC234 D - Prefix K-th Max

ふつうにheapqでやりましょう...

Submission #53027482 - AtCoder Beginner Contest 234