【SortedList】ABC281 - E - Least Elements
未履修のデータ構造についてアウトプット。
ABC281 E問題: Least Elements (diff 1334)
提出したコード
https://atcoder.jp/contests/abc281/submissions/52265648
整数$${M, K}$$と長さ$${N}$$の整数列$${A}$$が与えられる。
$${1 \leq i \leq N - M + 1}$$である全ての$${i}$$について$${A[i:i+M]}$$の小さい方から$${K}$$個の値の総和を出力してね。
SortedList とは
（私は中身まで理解していませんが）その名の通り「ソート順を保って要素の追加や削除ができるリスト」です。
Pythonには、c++の std::set や std::multiset に代わるものが標準で搭載されておらず、 それらが想定解の問題において難易度が跳ね上がっていました。
ABC-217 Cutting Woods が有名ですね。
それを解決してくれる sortedcontainers というライブラリが、AtCoder の2023年8月の言語アップデートで使えるようになっていたようです。（知らなかった・・・。）
具体的な使い方、および計算量については下記の記事をご参照ください。
使って解いてみる
$${A[i:i+M]}$$の$${M}$$個の要素の集合のうち、値の小さい方から$${K}$$個の要素の集合を$${S}$$、残りの要素の集合を$${T}$$とします。
$${i}$$が +1 されるとき、集合$${S, T}$$のどちらかから$${A_{i}}$$が削除、どちらかに$${A_{i + M}}$$が追加されます。
$${S}$$の要素の変化を丁寧に追いましょう。
はじめに$${A}$$の先頭$${M}$$個の要素について答えを求めておき、差分計算をすることで解くことができます。
$${S[-1]}$$は$${S}$$の最大値、$${T[0]}$$は$${T}$$の最小値を表します。
from sortedcontainers import SortedList
N, M, K = map(int, input().split())
A = list(map(int, input().split()))
""" i = 0 の答えを求める """
B = sorted(A[:M])
S = SortedList(B[:K])
T = SortedList(B[K:M])
res = sum(S)
ans = [res]
""" i > 0 の答えを求める """
for i in range(N - M):
T.add(A[i + M])
if A[i] <= S[-1]:
t = T.pop(0)
S.add(t)
S.remove(A[i])
res -= A[i]
res += t
else:
T.remove(A[i])
if S[-1] > T[0]:
s, t = S.pop(), T.pop(0)
S.add(t)
T.add(s)
res += t - s
ans.append(res)
print(*ans)順序付き集合がないのは競プロにおける Python の弱点でもあるので、これを機に使いこなせるようになっておきたいです。