Day 108-Q1 | Zip Tree（ジップ木）

問題

整数の集合 $S$（初期状態は空）に対して、$Q$ 個のクエリを順に処理せよ。

1 x: $S$ に $x$ を追加する（$x$ はまだ $S$ に存在しないことが保証される）
2 x: $S$ から $x$ を削除する（$x$ は $S$ に存在することが保証される）
3 k: $S$ を昇順に並べたとき $k$ 番目（1-indexed）に小さい要素を出力する（$1 \le k \le |S|$ が保証される）
4 x: $S$ の要素のうち $x$ 未満のものの個数を出力する

Zip Tree（ランクに基づく乱択平衡二分探索木、Tarjan & Zhan, 2018）を用いて、挿入・削除を期待 $O(\log N)$、3/4 クエリを $O(\log N)$ で処理せよ。

入力形式

Q
query_1
query_2
...
query_Q

制約

$1 \le Q \le 2\times10^5$

$0 \le x \le 10^9$

クエリの整合性は保証される

入出力例

入力例1

出力例1

5,3,8を追加すると$S=\{3,5,8\}$。2番目に小さいのは5。6未満は$\{3,5\}$の2個。5を削除すると$S=\{3,8\}$。2番目に小さいのは8。6未満は$\{3\}$の1個。

概念図: unzip（挿入）と zip（削除）

ヒント（段階的開示）

ヒント1: 方向性

配列をbisectでソート済みに保つ実装は挿入・削除のたびに$O(N)$の要素シフトが発生する。順序統計（k番目、rank）を高速に扱うには、各ノードに部分木サイズを持たせた平衡二分探索木が必要になる。Treap（優先度をランダムに割り当てて期待的にバランスさせる木）はよく知られているが、Zip TreeはTreapと同じ「ランダムなランク」の発想を使いながら、挿入を「ランクに沿ってツリーを2本に分割して新ノードの周りに再結合する（unzip）」、削除を「2つの子部分木をランク順に交互に併合する（zip）」という操作で捉え直したデータ構造である。

ヒント2: アプローチ

各ノードにランダムなランク（整数、大きいほど「根に近い」）を割り当てる。ツリーは「キーについては通常の二分探索木順序」「ランクについては最大ヒープ順序（親のランク ≥ 子のランク）」の両方を満たす。

挿入（unzip）: 新しいキー$x$を挿入する位置を探すのではなく、木全体を「キーが$x$未満の部分」と「$x$以上の部分」の2本の木に分割する。分割後、新ノード$x$（ランダムなランク付き）を根とし、2本の木をそれぞれ左右の子としてランクに従って併合すればよい。
削除（zip）: 削除したいノード$t$が見つかったら、$t$の左部分木と右部分木を「ランクの大きい方を上に」という規則で1本の木に併合（zip）し、$t$の位置に差し替える。

いずれも、キーによる分割（split）とランクによる併合（merge）という2つの基本操作の組み合わせで実現でき、これは通常のTreapのsplit/merge実装と操作的に等価である。

ヒント3: 誘導（コード骨格）

def unzip(t, k):
    # tを「キー<kの木」「キー>=kの木」に分割して(L, R)を返す
    if t == 0:
        return 0, 0
    if key[t] < k:
        L, R = unzip(right[t], k)
        right[t] = L
        update_size(t)
        return t, R
    else:
        L, R = unzip(left[t], k)
        left[t] = R
        update_size(t)
        return L, t

def zip_(L, R):
    # キーが「Lの全要素 < Rの全要素」を満たす2本の木を、ランクの最大ヒープ性を保って併合
    if L == 0: return R
    if R == 0: return L
    if rank[L] > rank[R]:
        right[L] = zip_(right[L], R)
        update_size(L)
        return L
    else:
        left[R] = zip_(L, left[R])
        update_size(R)
        return R
# 挿入 = unzipしてnodeを挟んでzip、削除 = 見つけたノードのleft/rightをzipして差し替え

模範解答 (Python)

import sys, random

def solve():
    data = sys.stdin.buffer.read().split()
    idx = 0
    Q = int(data[idx]); idx += 1

    rnd = random.Random(12345)

    key = [0]
    rank = [0]
    left = [0]
    right = [0]
    size = [0]

    def new_node(k):
        key.append(k)
        rank.append(rnd.getrandbits(30))
        left.append(0)
        right.append(0)
        size.append(1)
        return len(key) - 1

    def upd(t):
        size[t] = 1 + size[left[t]] + size[right[t]]

    def unzip(t, k):
        if t == 0:
            return 0, 0
        if key[t] < k:
            L, R = unzip(right[t], k)
            right[t] = L
            upd(t)
            return t, R
        else:
            L, R = unzip(left[t], k)
            left[t] = R
            upd(t)
            return L, t

    def zip_(L, R):
        if L == 0: return R
        if R == 0: return L
        if rank[L] > rank[R]:
            right[L] = zip_(right[L], R)
            upd(L)
            return L
        else:
            left[R] = zip_(L, left[R])
            upd(R)
            return R

    root = 0

    def insert(k):
        nonlocal root
        L, R = unzip(root, k)
        node = new_node(k)
        root = zip_(zip_(L, node), R)

    def erase(t, k):
        if t == 0:
            return 0
        if key[t] == k:
            return zip_(left[t], right[t])
        elif k < key[t]:
            left[t] = erase(left[t], k)
        else:
            right[t] = erase(right[t], k)
        upd(t)
        return t

    def delete(k):
        nonlocal root
        root = erase(root, k)

    def kth(t, k):
        if size[left[t]] >= k:
            return kth(left[t], k)
        elif size[left[t]] + 1 == k:
            return key[t]
        else:
            return kth(right[t], k - size[left[t]] - 1)

    def count_less(t, x):
        if t == 0:
            return 0
        if key[t] < x:
            return size[left[t]] + 1 + count_less(right[t], x)
        else:
            return count_less(left[t], x)

    out = []
    for _ in range(Q):
        t = data[idx]; idx += 1
        x = int(data[idx]); idx += 1
        if t == b'1':
            insert(x)
        elif t == b'2':
            delete(x)
        elif t == b'3':
            out.append(str(kth(root, x)))
        else:
            out.append(str(count_less(root, x)))
    print("\n".join(out))


solve()

計算量: 各操作は期待 $O(\log N)$（ランダムランクによりZip Treeの期待高さが$O(\log N)$に収まることに依存）。20,000回のランダム操作をbisectベースの愚直実装と突き合わせるストレステストで一致を確認済み。

Step-by-Step 解説

1ノードを配列で表現する
Pythonではノードをオブジェクトにすると属性アクセスが遅いため、key, rank, left, right, sizeを並行配列（インデックス0を「空」を表す番兵として使う）で管理する。

2unzipでキーによる分割を行う
unzip(t, k)は木$t$を「キー$<k$」「キー$\ge k$」の2本に分割する。通常のTreapのsplitと全く同じ実装。

3zip_でランクによる併合を行う
zip_(L, R)は「$L$の全要素$<R$の全要素」という前提のもとで2本の木を1本にする。ランクが大きい方を根に残す。ランダムなランクのおかげで期待的に木の高さが$O(\log N)$に保たれる。

4挿入は分割→挟んで併合、削除は該当ノードをzipで置き換え
挿入はunzip(root,k)で分割後、新ノードを挟んでzip_で結合。削除は対象ノードの左右の子をzip_(left[t], right[t])で1本にまとめて置き換える。この「削除時に左右の子をzipする」操作がZip Treeの名前の由来。

5順序統計は部分木サイズで計算する
kthとcount_lessは各ノードが保持するsizeを使って$O(\log N)$で計算する。

よくあるミス

ミス	原因	正しい書き方
`unzip`/`zip_`の後に`upd(t)`を呼び忘れる	子を付け替えたのに`size`を更新しない	左右の子を変更した直後に必ず`upd(t)`を呼ぶ
挿入時に`zip_(node, R)`を先に呼ぶ順序を間違える	`zip_(L, node)`と`zip_(node, R)`はどちらも数学的には可換ではない（引数の順序＝キーの大小関係の前提）	`zip_(zip_(L, node), R)`（左から右へキーが昇順になる順序を厳守）
ランクの同点判定で`>=`を使い常に左側を優先してしまう	ランクが衝突したときに片方に偏り、期待計算量の保証が崩れる	`zip_`の分岐は`rank[L] > rank[R]`のように厳密不等号にする
削除時に対象ノードが見つからないまま再帰が底まで進む	削除保証がない入力に対してサイレントに失敗する	本問では削除対象の存在が保証されるため問題ないが、汎用実装では`t==0`時の扱いを明示する

次のステップ

発展: 多重集合（同じキーの重複を許す）に拡張し、count, erase_allなどのクエリを追加する
発展: 区間反転・区間和などの遅延伝播をZip Treeに組み込む（Treapと同様にlazyフィールドを追加すればよい）
次回予告: Block-Cut Tree（二重連結成分分解の木構造化・LCAによる同一ブロック判定）

自己評価

完全理解おおむね理解要復習

Day 108-Q1 — Zip Tree（ジップ木）

問題