Day 118-Q1 | AVL木（自己平衡二分探索木・順序統計クエリ）

問題

N個のクエリを順に処理せよ。多重集合（初期状態は空）に対し、次の2種類のクエリが与えられる。

I x: 値 x を多重集合に挿入する
K k: 現在の多重集合の中で k 番目に小さい値（1-indexed）を出力する。要素数が k 未満の場合は -1 を出力する。

入力形式

N
query_1
query_2
...
query_N

各 query_i は I x または K k の形式。

制約

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

$1 \le x \le 10^9$（重複あり）

$1 \le k \le N$

入出力例

入力例1

6
I 5
I 3
I 8
K 2
I 1
K 1

出力例1

5
1

I 5,I 3,I 8後の多重集合は{3,5,8}で2番目に小さいのは5。さらにI 1後は{1,3,5,8}で1番目に小さいのは1。

概念図: 挿入と回転によるバランス回復

ヒント（段階的開示）

ヒント1: 方向性

挿入とk番目クエリを両方O(log N)で処理するには、単純なソート済み配列への挿入（O(N)）や、何も工夫しない二分探索木（片方向に伸びると最悪O(N)）では間に合わない。木の高さを常にO(log N)に保証する自己平衡機構が必要で、さらに各ノードに部分木サイズを持たせれば順序統計（k番目）クエリも木の高さに比例する時間で解ける。

ヒント2: アプローチ

AVL木は各ノードの左右部分木の高さの差（平衡係数）を常に-1, 0, +1のいずれかに保つ。挿入後に平衡係数が±2になったノードで、崩れ方に応じて4パターン（LL/RR/LR/RL）のいずれかの回転を行うことで平衡を回復する。同時に各ノードにsize（部分木のノード数）を保持し、回転のたびに再計算すれば、k番目クエリは「左部分木サイズと比較しながら降りていく」だけでO(log N)になる。

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

# ノードは配列(left[], right[], height[], size[], val[])で管理し、0を「存在しない」に使う
# 挿入: 通常のBST挿入を再帰で行い、戻りながら height/size を更新し、
#       平衡係数が2以上/-2以下になったノードで回転する
#
# balance = height[left] - height[right]
# balance > 1  かつ left側の子がさらに左に偏っている(LL) -> 右回転のみ
# balance > 1  かつ left側の子が右に偏っている(LR)      -> 左の子を左回転してからLLと同じ処理
# balance < -1 の場合は対称(RR, RL)
#
# k番目クエリ: node を root から辿り、
#   left_size = size[left[node]]
#   k <= left_size なら左へ、k == left_size+1 ならその場のvalが答え、
#   それ以外は k -= left_size+1 として右へ

模範解答 (Python)

import sys


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

    left = [0] * (N + 1)
    right = [0] * (N + 1)
    height = [0] * (N + 1)
    size = [0] * (N + 1)
    val = [0] * (N + 1)
    cnt = 0
    root = 0

    def upd(x):
        l, r = left[x], right[x]
        height[x] = max(height[l], height[r]) + 1
        size[x] = size[l] + size[r] + 1

    def bf(x):
        return height[left[x]] - height[right[x]]

    def rotate_right(y):
        x = left[y]
        t2 = right[x]
        right[x] = y
        left[y] = t2
        upd(y); upd(x)
        return x

    def rotate_left(x):
        y = right[x]
        t2 = left[y]
        left[y] = x
        right[x] = t2
        upd(x); upd(y)
        return y

    def insert(node, v):
        nonlocal cnt
        if node == 0:
            cnt += 1
            val[cnt] = v
            left[cnt] = right[cnt] = 0
            height[cnt] = 1
            size[cnt] = 1
            return cnt
        if v < val[node]:
            left[node] = insert(left[node], v)
        else:
            right[node] = insert(right[node], v)
        upd(node)
        b = bf(node)
        if b > 1:
            if bf(left[node]) < 0:
                left[node] = rotate_left(left[node])
            return rotate_right(node)
        if b < -1:
            if bf(right[node]) > 0:
                right[node] = rotate_right(right[node])
            return rotate_left(node)
        return node

    def kth(node, k):
        while node:
            ls = size[left[node]]
            if k <= ls:
                node = left[node]
            elif k == ls + 1:
                return val[node]
            else:
                k -= ls + 1
                node = right[node]
        return -1

    out = []
    for _ in range(N):
        t = data[idx]; idx += 1
        if t == b'I':
            x = int(data[idx]); idx += 1
            root = insert(root, x)
        else:
            k = int(data[idx]); idx += 1
            out.append(str(kth(root, k)))
    print('\n'.join(out))


sys.setrecursionlimit(1 << 20)
solve()

計算量: 挿入・k番目クエリともに木の高さに比例し、AVL木の高さは常に $O(\log N)$ に保たれるため、全体で $O(N\log N)$。乱数200試行のstress test（brute forceのbisect.insortによるソート済み配列と突き合わせ）で一致を確認済み、かつ全挿入後に平衡条件 $|\text{balance}|\le 1$ が破れていないことも検証済み。

Step-by-Step 解説

1なぜ単純なBSTでは不十分か
何も工夫しない二分探索木は、挿入順序が偏る（昇順・降順に近い）と片方向に伸び、高さがO(N)になり最悪ケースでクエリ1回がO(N)になってしまう。

2平衡係数と回転
balance = height[left]-height[right]。$|\text{balance}|>1$になったノードで回転して解消する。LL/RR/LR/RLの4パターンのうち、崩れた方向と子の偏り方の組み合わせで判定する。

3部分木サイズを使った順序統計
各ノードにsize（自分を含む部分木のノード数）を保持しておけば、「左部分木に何個あるか」を見ながら降りるだけでk番目要素が求まる。

4回転時のsize/heightの更新順序
回転で親子関係が入れ替わるので、必ず「新しい子(元の親)→新しい親」の順でupd()を呼ぶ。逆順だと新しい親の情報が古いまま計算されてしまう。

5計算量の確認
挿入・クエリともに$O(\log N)$、全体で$O(N\log N)$。$N\le2\times10^5$でも十分高速。

よくあるミス

ミス	原因	正しい書き方
回転後にsize/heightの更新順序を誤る	新しい親を先に更新してしまう	必ず「新しい子(元の親)→新しい親」の順でupd()する
bf(node)を古い子の状態で判定してしまう	挿入直後の再帰の戻り値を反映する前に平衡係数を見てしまう	upd(node)で自分の情報を更新してからbf(node)を評価する
kth関数でk==left_size+1の分岐を忘れる	「k<=lsなら左」「そうでなければ右」の2分岐だけにしてしまう	自分自身が「ls+1番目」であるケースを明示的に扱う
削除も含めて非再帰の反復処理だけで済まそうとする	AVLは回転で親子関係が変わるため片方向の走査だけでは対応しづらい	素直に再帰（または明示スタック）で実装する

次のステップ

発展: 削除操作（delete）を追加し、削除後も平衡が保たれることを確認する。
次回予告: 強力数法（Powerful Number Method）

自己評価

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

Day 118-Q1 — AVL木（自己平衡二分探索木・順序統計クエリ）

問題