Day 152-Q4 | 永続AVL木（Persistent AVL Tree）

問題

はじめ、空の集合をバージョン $0$ とする。以下 $Q$ 個のクエリを順に処理せよ。クエリは2種類ある。

1 v : 現在の最新バージョンに整数 $v$ を1つ挿入し、それを新しいバージョンとする（バージョン番号は挿入クエリが来るたびに $1$ ずつ増える）。
2 t k : バージョン $t$ の集合の中で $k$ 番目に小さい値を出力する（$1$-indexed）。バージョン $t$ の要素数が $k$ 未満、または $k<1$ の場合は -1 を出力する。

集合は多重集合（同じ値が複数回挿入されうる）として扱う。挿入は常に最新バージョンに対して行う（過去バージョンからの分岐は無い簡略版）。

入力形式

Q
query_1
...
query_Q

制約

$1 \le Q \le 200000$

1 vの$v$は$1 \le v \le 10^9$

2 t kの$t$はその時点で存在するバージョン番号

入出力例

入力例1

出力例1

入力例2

出力例2

10
20
10

概念図：パスコピーによる永続化（入力例2）

ヒント（段階的開示）

ヒント1（方向性）

通常のAVL木は挿入のたびに既存ノードを直接書き換える。永続化するには、書き換えが必要なノードを新しいオブジェクトとして作り直し、書き換えが不要な子は元のノードをそのまま参照する。「バージョンごとの根」を配列に保存しておくだけで過去の状態が全部残る。

ヒント2（アプローチ）

挿入処理を再帰で書き、値を挿入する側の子だけを新しく作って返す（もう片方の子は元のオブジェクトをそのまま使う）。回転処理も同様に、位置が変わるノードだけ新規作成する。$k$番目探索は各ノードに部分木サイズを持たせ、左部分木のサイズと比較しながら降りる。

ヒント3（誘導）

class Node:
    __slots__ = ('val', 'left', 'right', 'height', 'size')
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
        self.height = 1 + max(_h(left), _h(right))
        self.size = 1 + _s(left) + _s(right)

def insert(n, v):
    if n is None:
        return Node(v)
    if v < n.val:
        new_n = Node(n.val, insert(n.left, v), n.right)
    else:
        new_n = Node(n.val, n.left, insert(n.right, v))
    return _rebalance(new_n)   # 回転が必要なら新しいノードを組み立てて返す

模範解答 (Python)

import sys


class Node:
    __slots__ = ('val', 'left', 'right', 'height', 'size')

    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
        self.height = 1 + max(_h(left), _h(right))
        self.size = 1 + _s(left) + _s(right)


def _h(n):
    return n.height if n else 0


def _s(n):
    return n.size if n else 0


def _balance(n):
    return _h(n.left) - _h(n.right) if n else 0


def _rotate_right(n):
    l = n.left
    new_n = Node(n.val, l.right, n.right)
    new_l = Node(l.val, l.left, new_n)
    return new_l


def _rotate_left(n):
    r = n.right
    new_n = Node(n.val, n.left, r.left)
    new_r = Node(r.val, new_n, r.right)
    return new_r


def _rebalance(n):
    b = _balance(n)
    if b > 1:
        if _balance(n.left) < 0:
            n = Node(n.val, _rotate_left(n.left), n.right)
        return _rotate_right(n)
    if b < -1:
        if _balance(n.right) > 0:
            n = Node(n.val, n.left, _rotate_right(n.right))
        return _rotate_left(n)
    return n


def insert(n, v):
    if n is None:
        return Node(v)
    if v < n.val:
        new_n = Node(n.val, insert(n.left, v), n.right)
    else:
        new_n = Node(n.val, n.left, insert(n.right, v))
    return _rebalance(new_n)


def kth(n, k):
    ls = _s(n.left)
    if k <= ls:
        return kth(n.left, k)
    elif k == ls + 1:
        return n.val
    else:
        return kth(n.right, k - ls - 1)


def main():
    data = sys.stdin.buffer.read().split()
    idx = 0
    q = int(data[idx]); idx += 1
    versions = [None]
    out = []
    for _ in range(q):
        typ = data[idx]; idx += 1
        if typ == b'1':
            v = int(data[idx]); idx += 1
            versions.append(insert(versions[-1], v))
        else:
            t = int(data[idx]); idx += 1
            k = int(data[idx]); idx += 1
            root = versions[t]
            size = _s(root)
            if k < 1 or k > size:
                out.append('-1')
            else:
                out.append(str(kth(root, k)))

    print('\n'.join(out))


main()

Step-by-Step 解説

1入力例1をバージョンごとに追う
version1={5}, version2={3,5}, version3={3,5,8}。2 3 2はバージョン3のソート済み{3,5,8}の2番目=5。

2入力例2で「永続性」を確認する
1 5でversion3={5,10,20}を新たに作った直後の2 2 1が、挿入前と全く同じ答え10を返すことが永続化の証拠。

3パスコピーで新規作成されるノード数を確認
1 5実行時、値5はversion2の根(10)より小さいので左へ。根と新しい葉(5)の2つだけ新規作成され、右の子(20)は既存オブジェクトを再利用する。挿入1回あたり$O(\log N)$個の新規ノードで済む。

4回転が発生しても「作り直すのはパス上だけ」
回転関数も新しいNodeを作って返すだけで、回転に関与しない部分木は既存オブジェクトをそのまま子として持たせる。

よくあるミス

ミス	原因	正しい書き方
挿入時に既存ノードのフィールドを直接書き換えてしまう	通常のAVL木の実装をそのまま流用	必ず新しい`Node(...)`を作って返し、古いノードには触れない
回転関数内で元の`height`/`size`を使い回す	再計算忘れ	`Node.__init__`で毎回子から再計算する設計にする
バージョン番号とクエリ全体の番号を混同する	インデックスの取り違え	`versions`は挿入クエリの回数分だけ伸びる配列
$k$番目探索で部分木を毎回全走査してしまう	順序統計木の考え方を使っていない	各ノードに部分木サイズを持たせ$O(\log N)$で降りる
過去バージョンにも挿入できると誤解する	問題の簡略化を読み落とす	本問では常に最新バージョンにのみ挿入する

次のステップ

発展: 「過去の任意のバージョンに対して挿入し新しいバージョンを分岐させる」機能を実装し、バージョン管理をDAG構造にしてみる。
次回予告: Q学習（Q-learning）

自己評価

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

Day 152-Q4 — 永続AVL木（Persistent AVL Tree）

問題