【SortedList】ABC281 - E - Least Elements

未履修のデータ構造についてアウトプット。

ABC281 E問題: Least Elements (diff 1334)

E - Least Elements

提出したコード
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月の言語アップデートで使えるようになっていたようです。（知らなかった・・・。）

具体的な使い方、および計算量については下記の記事をご参照ください。

そろそろ我々はsortedcontainersを使えるようになった方がいい - Qiita

[競プロ python] ふんわりと理解！SortedListの仕組み

使って解いてみる

$${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 の弱点でもあるので、これを機に使いこなせるようになっておきたいです。