Python で std::set の代替物を作った

Pythonista の皆様、こんにちは。競技プログラミングをしている tatyam と申します。

Python で競技プログラミングをしていると、std::set が欲しいな〜ってとき、ありますよね？
例えばそう、

Cutting Woods (ABC217-D)

を見てみましょう。これは以下のような問題です。

長さ$L$の定規がある。以下のクエリが$Q$回与えられるので、順に処理せよ。
1 x：目盛$x$で定規を切る。
2 x：目盛$x$を含む部分の長さを出力する。
$L\le {10}^9,N\le 2\times {10}^5$

素直に考えると、切れ目の位置の集合を管理すれば良さそうです。

• $s$を切れ目の位置の集合とする。はじめ、$s=0,L$である。
• 1 x が与えられたら、$s$に$x$を追加する。
• 2 x が与えられたら、($s$の$x$以上の要素のうち最小 )$-$($s$の$x$以下の要素のうち最大 ) を出力する。

といったアルゴリズムが考えられます。あとは、「集合に要素を追加する」「$x$以下で最小の要素を取得する」操作が高速にできるデータ構造があればいいですね。

それを実現するのが 平衡二分探索木 (C++ における std::set) というわけです。

平衡二分探索木

平衡二分探索木は、

• 要素の追加
• 要素の削除
• $x$以上の最小の要素の検索

などの操作が要素数を$N$として$O(\log N)$の計算量で行えるすごいデータ構造です。

C++ であれば set 、Java であれば TreeSet 、C# であれば SortedSet と、多くの言語で標準ライブラリに入っています。Cutting Woods の公式解説も平衡二分探索木を使ったものになっています。
しかし、Python の標準ライブラリには平衡二分探索木がありません！！

外部ライブラリは使うことができず、コピペにも長すぎ、
C++ でバイナリを作るのはめちゃくちゃ速いけどちょっと長いしどこでも使えるものではなく…
平衡二分探索木が使えないとなると、クエリを逆向きにして UnionFind で長さを管理するなど、考察も実装も一段と難しいものになってしまいます。

そこで、

• コピペして提出できる
• 実装が短い
• 速い

std::set の代替物を作りましょう。

実装方針

std::set の実現方法にもいろいろあって、

汎用的なやつ

• 平衡二分探索木 (AVL 木、赤黒木、Treap、Splay 木 など)

整数しか扱えないやつ

• BinaryTrie
• https://qiita.com/Kiri8128/items/6256f8559f0026485d90
• 64分木 (BinaryTrie を 64 分木にして bit 演算で高速化したやつ)

などがあります。
今回は、風の噂で平方分割が速いと聞いたので、平方分割を試してみましょう。

平方分割

平方分割では、要素を$\sqrt{N}$個ずつ別々の list に入れ、list の list の形で要素を管理します。
例えば、$\{1,3,4,5,8,9,11,14,16\}$という集合は [[1, 3, 4], [5, 8, 9], [11, 14, 16]] という形です。
このとき、$\sqrt{N}$個できた小さな list を bucket と呼ぶことにします。bucket の数は$\sqrt{N}$個、各 bucket に入っている要素数も$\sqrt{N}$個です。

このように管理すると、「$x$以上の最小の要素の検索」は効率的にできることがわかります。

1. bucket を順に見ていって$x$以上の要素が含まれるか調べる
2. $x$以上の要素が含まれる最初の bucket について、$x$以上であるような最初の要素を見つける

とすれば$O(\sqrt{N})$ですね。 (二分探索すれば$O(\log N)$になります。)

「要素の追加・削除」についてはどうでしょうか？

要素の追加については、適切な bucket の適切な位置に insert 、要素の削除については、$O(\sqrt{N})$でどこにあるか見つけて pop とすれば$O(\sqrt{N})$になりそうですが、これを何度もやって bucket の大きさが変わると$O(\sqrt{N})$とは行かなくなってきます。
そこで、要素の追加・削除を$\sqrt{N}$回行ったところで bucket を壊し、再び$\sqrt{N}$個ずつ新しい bucket に振り分けるようにしてみましょう。
この再構築には$O(N)$の計算量が掛かりますが、再構築するのには$\sqrt{N}$回の操作が必要なので、amortized$O(\sqrt{N})$の計算量になります。

平方分割の特徴

• 汎用的

平衡二分探索木と同様に、整数以外でも比較可能であれば何でも入れることができます。

• 速い

平衡二分探索木は各操作が$O(\log N)$、平方分割は各操作$O(\sqrt{N})$であるため、平衡二分探索木の方が速そうに見えますが、うまく高速化すれば、(競プロ実用の範囲では) 平衡二分探索木と同等かそれ以上の速さを得ることができます。
これには、Python の言語特性だったり、list という組み込み型を使えることが影響していそうです。

• 多機能

C++ の std::set は、

• 小さい方から$x$番目の要素は何か？
• $x$は小さい方から何番目か？

といったクエリに答えてくれる関数がありません。これを効率的に行うには、平衡二分探索木の各ノードに、「自分より下にいくつノードがあるか」を持たなければならず、メモリ使用量が多くなってしまいます。(これをする平衡二分探索木も用意されてはいます。)

平方分割なら、各 bucket は list なので、bucket の中の要素数を数えることは簡単です。そのため、「小さい方から$x$番目の要素は何か？」「$x$は小さい方から何番目か？」といったクエリに答えることができます。

• 実装が短い

平衡二分探索木と比べて、平方分割は list の list という単純な構造なので、実装も短くなります。

SortedSet

上に書いたような、平方分割による std::set の代替物にいい感じの機能をつけて高速化したものがこちらになります。

高速化のために、

• (bucket の個数) : (各 bucket に入っている要素数) = 1 : 50 程度になるように分割
• 1 : 50 が 1 : 0 または 1 : 170 になったところで bucket を再構築

という実装になっていることに注意してください。

使ってみよう

Cutting Woods (ABC217-D)

試しに、この SortedSet を使って Cutting Woods (ABC217-D) を解いてみましょう。

s = SortedSet([0, L]) で$0$と$L$が入った集合を作り、
s.add(x) で$x$を追加し、
s.gt(x) - s.lt(x) で答えを求められます。

# paste SortedSet here

L, Q = map(int, input().split())
s = SortedSet([0, L])
for i in range(Q):
    c, x = map(int, input().split())
    if c == 1:
        s.add(x)
    else:
        print(s.gt(x) - s.lt(x))

短い！！そして、377 ms はなかなかの速さです。

これは細かいポイントですが、

print(s.gt(x) - s.lt(x))

を

i = s.index(x)
print(s[i] - s[i - 1])

と書くと若干速くなります。
(2 回 bucket の中を探索していたのが 1 回に減るので)

さらに、

a = s.a
for i in range(len(a)):
    if a[i][-1] > x:
        j = bisect_left(a[i], x)
        print(a[i][j] - (a[i][j - 1] if j else a[i - 1][-1]))
        break

とベタ書きすると一番速いです。
(どの bucket にあるかを探すのが 1 回に減るので)

このような場合に使えるように、内部の index を取れるように改造するのも良いでしょう。

Prefix K-th Max (ABC234-D)

ヒープを使うのが想定のこの問題も、添字アクセスができる SortedSet で殴ることができます。

# paste SortedSet here

N, K = map(int, input().split())
P = list(map(int, input().split()))
K -= 1
s = SortedSet(P[:K])
for i in range(K, N):
    s.add(P[i])
    print(s[~K])

おわり

SortedSet をぜひご活用ください！