Pythonのsortedcontainersのチートシートです。
まえがき
競技プログラミングをやっている人向け:
長年AtCoderのPyPy環境にはC++でいうstd::setやstd::mapにあたるデータ構造が標準で実装されておらず、それらが想定解で要求される問題ではより難しい別解を実装したり、外部のライブラリを借りてきたりする必要がありました。この度、AtCoderのジャッジアップデートによってsortedcontainersライブラリが使えるようになりました。これは非常に強力なデータ構造ライブラリなので、これを機に勉強して使いこなせるようになっておきましょう。
競技プログラミングをやっていない人向け:
sortedcontainersライブラリはPython用のデータ構造ライブラリです。本記事では、https://grantjenks.com/docs/sortedcontainers/# をもとに、SortedSet, SortedList, SortedDict それぞれについてできることを解説し、操作の使い方や計算量について紹介します。
はじめに
何はともあれ、手元環境で使えるようにするためにインストールしましょう。コンソールでpipを用いてインストールできます。
$ pip install sortedcontainers
また、Pythonコードで sortedcontainers を import するわけですが、私は
from sortedcontainers import SortedSet, SortedList, SortedDict
とすることをお勧めします。sortedcontainers.sortedset(大文字小文字が異なる)などを使うと、意図した動作になりません。↑を書いておけば、次から説明するコード例がそのまま動きます。
以降、特に言及がなければNはその瞬間に対象のデータ構造に格納されている要素数であるとします。
SortedSet
簡単に言うと set の各要素が常にソートされているデータ構造だと思ってよいです。
初期化
S = SortedSet([3, 1, 2])
#SortedSet([1, 2, 3]), 初期化の計算量は O(N * logN)
S.add(4)
#SortedSet([1, 2, 3, 4]), add の計算量は O(logN)
S.add(3)
#SortedSet([1, 2, 3, 4]), 要素は重複しない
S.discard(4)
#SortedSet([1, 2, 3]), 値4を削除 計算量は O(logN) 存在しない要素をdiscardすると、何も起こらない
#S.remove(100) とやると、KeyErrorが出る
S.pop(2)
#S[2]を削除 S.pop()で最大要素の削除 S.pop(0) で最小要素の削除 全部 O(logN)
S[1]
#2 getは O(logN)
S[0:2]
#0 1 スライスなどもできる これの計算量は長さを k として O(klogN)
len(S)
#2 現在の要素数 O(1)
S.bisect_left(1)
#1
S.bisect_right(1)
#2
#これらは二分探索。 S.bisect_left(x) で、Sにxを挿入する位置(index)を返す。
#S.bisect_right(x) との違いは、すでにxがSにあるときに左に入れるか右に入れるか
#Pythonのbisectと同じ使用感
S.index(2)
#1 Sに2が現れる最初の位置を返す。ないとValueError
S.irange(0,2)
#[1, 2] S に含まれる 0以上2以下(両端含む)の要素を列挙
SortedList
ほとんどSortedSetと同じですが、重複した要素が許されます。
S = SortedList([3, 1, 2, 1])
#SortedList([1, 1, 2, 3]), 初期化の計算量は O(N * logN)
S.add(4)
#SortedList([1, 1, 2, 3, 4]), add の計算量は O(logN)
S.add(3)
#SortedList([1, 1, 2, 3, 3, 4]), 要素は重複する
S.discard(3)
#SortedList([1, 1, 2, 3, 4]), 値3を削除 計算量は O(logN) 要素は1個だけ消される
#S.remove(100) とやると、KeyErrorが出る
S.pop(2)
#S[2]を削除 S.pop()で最大要素の削除 S.pop(0) で最小要素の削除 全部 O(logN)
S[1]
#2 getは O(logN)
S[0:2]
#スライスなどもできる これの計算量はO(klogN)
len(S)
#現在の要素数 O(1)
S.bisect_left(1)
S.bisect_right(1)
#これらは二分探索。 S.bisect_left(x) で、Sにxを挿入する位置(index)を返す。
#S.bisect_right(x) との違いは、すでにxがSにあるときに左に入れるか右に入れるか
#Pythonのbisectと同じ使用感
S.index(2)
#1 Sに2が現れる最初の位置を返す。ないとValueError
S.count(1)
#2 Sに含まれる要素の個数を返す SortedSetにもあるがそっちは使う意味ないので…
S.irange(0,2)
#[1, 2] S に含まれる 0以上2以下(両端含む)の要素を列挙
SortedDict
Dictの要素がソートされたバージョンです。
S = SortedDict()
S['c'] = 2
S['b'] = 3
S['a'] = 1
#SortedDict({'a': 1, 'b': 3, 'c': 2}) keyの昇順に格納される O(logN)
S.items()
S.keys()
S.values()
#dictと同じ
S.pop()
#Sのうち、keyが最大のものを返す S.pop(i)でi番目 S.pop(0)で最小 全部 O(logN)
'c' in S
#True
使用例
AtCoderの問題 ABC217-D 「Cutting Woods」を解いてみるとこうなります。コード量が少ないですね。
https://atcoder.jp/contests/abc217/submissions/47087829
注意点
sortedcontainersの利用において,max関数、min関数を使用すべきではありません。これらの関数ではiterableを取得してmax,minをとるので,計算量が $\Omega (N)$ となってしまいます。
宣伝(本質情報)
tatyamさんが制作したSortedSetの実装があります。雑に比較したところsortedcontainersよりこちらの方が速いケースが多かったです(そ、そんな…)
https://github.com/tatyam-prime/SortedSet