【Python】SortedSet, SortedDictチートシート（使い方）

tadanori

pythonのset, dictはソートされていないため、C++のsetなどを前提としたAtCoderの問題では毎回実装を考える必要がありました。最近、sortedcontainersがAtCoderのジャッジでも利用できるようになったようなので、ソート付きのSortedSet, SortedList, SortedDictについて解説したいと思います。

その他のAtCoderに役立つ記事の一覧は以下にあります。

はじめに

AtCoderなどの競技プログラミングに参加していると、C++のsetのようなsorted setが標準で用意されていないため少し工夫が必要でした。

「SortedSetがあれば楽勝なのに」という問題に直面すると「不利だなぁ」と思うことも多かったと思います。

最近、他人のコードを読んでいて気づいたのですが、SortedSetをAtCoderのジャッジがサポートしたようです。

これを機に、利用できるようになったsortedcontainersのSortedSetの使い方について整理してみました。

SotedSet

わかりやすく説明すると集合型（set）で、要素がソートされているデータ構造です。

ソートされていることと、bisect_left, bisect_rightが使える点が大きいです。

SortedSetが使えたら、実装簡単なのにと思うことが度々ありました。これが使えると実装がかなり楽になります

from sortedcontainers import SortedSet

s = SortedSet([2, 3, 1, 44, 5, 4])
print(s)
#SortedSet([1, 2, 3, 4, 5, 44])　初期化
s.add(10)
print(s)
#SortedSet([1, 2, 3, 4, 5, 44])　addで追加
s.add(2)
print(s)
#SortedSet([1, 2, 3, 4, 5, 10, 44])　重複した要素は１つになる
print(s[0])
#1 s[0]で最も小さいものを選択
print(s[-1]) 
#44 s[-1]で最も小さいものを選択
s.pop()
print(s)
#SortedSet([1, 2, 3, 4, 5, 10]) 一番大きな要素が削除される
print(s.index(2))
#1
print(s.bisect_left(5)) 二分探索. C++のlower_bound()相当
#4
print(s.bisect_right(5)) 二分探索. C++のupper_bound()相当
#5
print(list(s.irange(3, 8)))
#[3,4,5] 範囲の値を取り出す
s.remove(4)
print(s)
#SortedSet([1, 2, 3, 5, 10]). 要素を削除する

SortedList

SortedSetと異なり重複した値の要素を取り扱うことができます。

小さい順、大きい順に取り出すのは優先キューでOKですが、bisect_left/bisect_rightによる二分探索ができるのは大きいです。

from sortedcontainers import SortedList

s = SortedList([3, 2, 1, 44, 5, 4])
print(s)
#SortedList([1, 2, 3, 4, 5, 44]) 　初期化
s.add(10)
print(s)
#SortedList([1, 2, 3, 4, 5, 10, 44]) addで追加
s.add(2)
print(s)
#SortedList([1, 2, 2, 3, 4, 5, 10, 44]) 重複した要素もOK
print(s[0])
#1 s[0]で最も小さいものを選択
print(s[-1])
#44 s[-1]で最も小さいものを選択
s.pop()
print(s)
#SortedList([1, 2, 2, 3, 4, 5, 10])  一番大きな要素が削除される
print(s.index(2))
#1
print(s.bisect_left(2)) 二分探索. C++のlower_bound()相当
#1
print(s.bisect_right(2)) 二分探索. C++のupper_bound()相当
#3
print(list(s.irange(3, 8)))
#[3,4,5] 範囲の値を取り出す
s.remove(2)
print(s)
#SortedSet([1, 2, 3, 5, 10]) 要素を削除する

SortedDict

dict（辞書）のソートされたものです。keyを取り出したときにソートされているのが特徴です。

from sortedcontainers import SortedDict

s = SortedDict({'d': 10})
print(s)
# SortedDict({'d': 10})
s['b'] = 2
s['e'] = 2
s['a'] = 2
print(s)
# SortedDict({'a': 2, 'b': 2, 'd': 10, 'e': 2})
for e in s.items() : print(e)
#('a', 2)
#('b', 2)
#('d', 10)
#('e', 2)
for e in s.keys() : print(e)
#a
#b
#d
#e
for e in s.values() : print(e)
#2
#2
#10
#2