AtCoderで使えるPythonのSortedSet, SortedDictを解説

Aru

Python標準のset, dictはソートされていないため、C++STLのsetを前提として競プロの問題では不利でした。実は、sortedcontainersがAtCoderのジャッジでも利用できるようになっているので、これを使えばC++と同様にソート付きのsetなどを使うことができます。この記事では、ソート付きのSortedSet, SortedList, SortedDictのサンプルコードを紹介します。

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

はじめに

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

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

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

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

なお、この記事ではプログラム中にコメントを入れる形で使い方を解説しています（こちらの方がイメージが掴みやすいかなと思い、この形式にしました）。

SotedSet

SortedSetは、わかりやすく説明すると集合型（set）で要素がソートされているものです。重複を排除して、かつ、ソートしておきたい場合にこのデータ構造を使います。

この、SortedSetを使うには、

from sortedcontainers import SortedSet

のインポート必要となります。

ソートされているため、先頭の要素は最も小さい値、末尾の要素は最も大きな値となります。また、２分探索による検索（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))
#4 二分探索. C++のlower_bound()相当

print(s.bisect_right(5))
#5 二分探索. C++のupper_bound()相当

print(list(s.irange(3, 8)))
#[3,4,5] 範囲の値を取り出す

s.remove(4)
print(s)
#SortedSet([1, 2, 3, 5, 10]). 要素を削除する

SortedList

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

なお、SortedListを使うには、

from sortedcontainers import SortedList

のインポート必要となります。

小さい順、大きい順に取り出すのは優先キュー（Pythonの場合はheap）で十分ですが、SortedListでは、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)) 
#1 二分探索. C++のlower_bound()相当

print(s.bisect_right(2)) 
#3 二分探索. C++のupper_bound()相当

print(list(s.irange(3, 8)))
#[3,4,5] 範囲の値を取り出す

s.remove(2)
print(s)
#SortedSet([1, 2, 3, 5, 10]) 要素を削除する

SortedDict

SortedDictは、辞書型（dict）でソートされているものになります。

SortedDictを使うには、

from sortedcontainers import SortedDict

のインポート必要となります。

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

使い道としては、SortedSet、SortedListよりは少ないかもしれません。

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