균형 이진트리(AVL/RB)

class Node:
    def __init__(self, key):
        self.key, self.left, self.right, self.h = key, None, None, 1

def h(n): return n.h if n else 0
def update(n): n.h = 1 + max(h(n.left), h(n.right))
def balance(n): return h(n.left) - h(n.right)    # 균형 인수

def rotate_right(z):
    y = z.left
    z.left, y.right = y.right, z
    update(z); update(y)
    return y                          # 서브트리의 새 루트

def rotate_left(z):
    y = z.right
    z.right, y.left = y.left, z
    update(z); update(y)
    return y

def insert(n, key):
    if not n: return Node(key)
    if key < n.key: n.left = insert(n.left, key)
    else: n.right = insert(n.right, key)
    update(n)
    b = balance(n)
    if b > 1:                                  # 왼쪽이 무거움
        if balance(n.left) < 0: n.left = rotate_left(n.left)    # LR: 먼저 왼쪽 회전
        return rotate_right(n)
    if b < -1:                                 # 오른쪽이 무거움
        if balance(n.right) > 0: n.right = rotate_right(n.right)  # RL: 먼저 오른쪽 회전
        return rotate_left(n)
    return n

root = None
for k in [10, 20, 30, 40, 50, 25]:
    root = insert(root, k)
print(root.key, root.left.key, root.right.key)   # 30 20 40

def bst_insert(n, key):          # 균형을 맞추지 않는 일반 삽입
    if not n: return Node(key)
    if key < n.key: n.left = bst_insert(n.left, key)
    else: n.right = bst_insert(n.right, key)
    update(n)
    return n

plain = avl = None
for k in range(1, 16):          # 1, 2, ..., 15를 정렬된 순서로
    plain = bst_insert(plain, k)
    avl = insert(avl, k)
print(h(plain), h(avl))        # 15 4

import java.util.TreeMap;

TreeMap<Integer, String> map = new TreeMap<>();
map.put(30, "c"); map.put(10, "a"); map.put(20, "b");

System.out.println(map.firstKey());      // 10  (가장 작은 키)
System.out.println(map.floorKey(25));    // 20  (25 이하 중 가장 큰 키)
System.out.println(map.ceilingKey(25));  // 30  (25 이상 중 가장 작은 키)
System.out.println(map.headMap(30));     // {10=a, 20=b}  (30 미만)