Day 113-Q1 — AA Tree（AA木・レベルベースの簡易平衡二分探索木による順序統計クエリ）

2026-08-05 赤色 Master / Phase 8+ ★★★★★★★★★ skew / split 操作による平衡・順序統計木

問題

初期状態で空の整数集合に対し、$Q$ 個のクエリを順番に処理せよ。クエリは次の3種類である。

  • 1 x — 集合に値 $x$ を追加する（同じ値が複数回追加されることはない）
  • 2 x — 集合の要素のうち $x$ 未満であるものの個数を出力する
  • 3 k — 集合の要素を昇順に並べたとき $k$ 番目（$1$-indexed）の値を出力する

入力形式

Q
query_1
query_2
...
query_Q

制約

$1 \le Q \le 2\times10^5$
クエリ1の $x$ は $1\le x\le10^9$ で相異なる
クエリ2の $x$ は $1\le x\le10^9$
クエリ3の $k$ はその時点の要素数以下

入出力例

入力例1

9
1 5
1 3
1 8
1 1
2 5
3 1
3 4
2 100
2 1

出力例1

2
1
8
4
0

集合は $\{1,3,5,8\}$ になる。5未満は$\{1,3\}$で2個。1番目に小さいのは1、4番目は8。100未満は全4個。1未満は0個。

概念図: skew（右回転）と split（左回転）で平衡を保つ

左の子と同レベル→skew／右の右の孫と同レベル→split skew（右回転）前 P(L1) L(L1) → skew後（Lが親になる） L(L1) P(L1) split（左回転）前: 右の右の孫が同レベル A(L2) B(L2) C(L2) → Bを新しい根にしてlevel(B)+=1 A.right = B.left, B.left = A 挿入後は skew → split の順で毎回適用するだけで木全体が O(log N) の高さに保たれる（削除なしの場合）

ヒント（段階的開示）

ヒント1: 方向性
挿入のたびに単純な二分探索木を使うと、挿入順によっては木が一直線に伸びてしまい、クエリ2・3が $O(N)$ になる。木の高さを $O(\log N)$ に保つ平衡二分探索木が必要。赤黒木は実装が複雑なので、「レベル」という単純な整数だけで平衡を管理するAA木（Arne Andersson考案の赤黒木の簡略版）を使う。
ヒント2: アプローチ
AA木は次の不変条件を保つ。①左の子のレベルは親より必ず1小さい ②右の子のレベルは親と同じか1小さく、かつ「右の右の孫」が親と同レベルになってはいけない。挿入後にこれが崩れたらskew（左の子が親と同レベルなら右回転）とsplit（右の右の孫が親と同レベルなら左回転してレベルを上げる）の2操作だけで修復できる。各ノードに部分木サイズを持たせれば順序統計が $O(\log N)$ で求まる。
ヒント3: 誘導（コード骨格） 
def insert(t, x):
    if t is NIL: return new_node(x)
    if x < val[t]: left[t] = insert(left[t], x)
    else: right[t] = insert(right[t], x)
    update_size(t)
    t = skew(t)
    t = split(t)
    return t

模範解答 (Python)

import sys
sys.setrecursionlimit(10000)

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

    NIL = 0
    val = [0]
    left = [0]
    right = [0]
    level = [0]
    size = [0]

    def new_node(x):
        val.append(x)
        left.append(NIL)
        right.append(NIL)
        level.append(1)
        size.append(1)
        return len(val) - 1

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

    def skew(t):
        if t == NIL or left[t] == NIL:
            return t
        if level[left[t]] == level[t]:
            l = left[t]
            left[t] = right[l]
            right[l] = t
            upd(t); upd(l)
            return l
        return t

    def split(t):
        if t == NIL or right[t] == NIL or right[right[t]] == NIL:
            return t
        if level[t] == level[right[right[t]]]:
            r = right[t]
            right[t] = left[r]
            left[r] = t
            level[r] += 1
            upd(t); upd(r)
            return r
        return t

    def insert(t, x):
        if t == NIL:
            return new_node(x)
        if x < val[t]:
            left[t] = insert(left[t], x)
        else:
            right[t] = insert(right[t], x)
        upd(t)
        t = skew(t)
        t = split(t)
        return t

    def rank_lt(t, x):
        res = 0
        while t != NIL:
            if val[t] < x:
                res += size[left[t]] + 1
                t = right[t]
            else:
                t = left[t]
        return res

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

    root = NIL
    out = []
    for _ in range(Q):
        t = data[idx]; idx += 1
        if t == b"1":
            x = int(data[idx]); idx += 1
            root = insert(root, x)
        elif t == b"2":
            x = int(data[idx]); idx += 1
            out.append(str(rank_lt(root, x)))
        else:
            k = int(data[idx]); idx += 1
            out.append(str(kth(root, k)))

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

main()
計算量: 挿入・クエリともに $O(\log N)$。$Q=2\times10^5$ 件のランダムクエリで bisect を使った参照実装と全出力が一致することを確認済み。

Step-by-Step 解説

1ノードを配列で表現する
オブジェクトではなく val, left, right, level, size の5本の配列でノードを表現し、ノード番号で参照する。NIL=0を番兵にする。
2skew — 左の子と同レベルなら右回転
「左の子は親より厳密に1小さいレベル」という条件が崩れたら右回転で修復する。
3split — 右の右の孫と同レベルなら左回転してレベルを上げる
同レベルのノードが横に3つ連続してはいけないという条件を、左回転+レベル+1で修復する。
4skew→splitの順で適用し続ける
挿入の再帰から戻るたびに毎回適用するだけで木全体が $O(\log N)$ の高さに保たれる。
5部分木サイズで順序統計を実現する
sizeを管理しておけば、$x$未満の個数も$k$番目の値も根から葉までの1本道で $O(\log N)$ で求まる。

よくあるミス

ミス原因正しい書き方
skew/splitの後にsizeを更新し忘れる回転で親子関係が変わっても古いsizeが残る内側から外側の順で必ずsizeを再計算する
splitの条件をright[t]だけで判定する「右の子」と「右の右の孫」を混同するright[right[t]]のレベルを見て判定する
rank_ltで<=を使いx自身を含めてしまう「未満」と「以下」を取り違えるval[t] < xの厳密な判定を使う
再帰深さの設定を怠るAA木は高さO(logN)なので通常問題ないが環境依存の上限がある念のためsetrecursionlimitを余裕を持って設定する

次のステップ

  • 発展: 削除クエリ 4 x を追加し、後継ノード入れ替え+decrease_level+skew/splitの再適用を実装する
  • 発展: Scapegoat Tree（Day111 Q2）・Treap（Day024 Q5）と挿入のみのワークロードで実測速度を比較する
  • 発展: kthクエリの結果を「中央値ストリーミング」に応用する

自己評価

自分の回答

気づき・メモ