트리 코드
노트의 개념을 Python 3로 옮긴 코드다. 표준 라이브러리만 쓰고, 파일 끝의 assert 예제가 그대로 실행된다. 복사해 python3 파일이름.py로 돌려 보자.
트리의 낱말
"""트리의 낱말: 부모 표로 준 트리에서 뿌리·잎·깊이·높이·부분 트리 크기를 구한다.
시간: 모든 함수 O(n) — 노드마다 한 번씩 본다
공간: O(n) — 자식 표와 재귀 호출 스택
"""
from collections import defaultdict
def children_of(edges: list[tuple[str, str]]) -> dict[str, list[str]]:
"""(부모, 자식) 목록을 부모 → 자식 목록 표로 바꾼다."""
children: dict[str, list[str]] = defaultdict(list)
for parent, child in edges:
children[parent].append(child)
return children
def find_root(edges: list[tuple[str, str]]) -> str:
"""부모가 없는 유일한 노드를 돌려준다."""
parents = {parent for parent, _ in edges}
childs = {child for _, child in edges}
roots = parents - childs
if len(roots) != 1:
raise ValueError(f"뿌리가 하나가 아니다: {sorted(roots)}")
return roots.pop()
def leaves(edges: list[tuple[str, str]]) -> list[str]:
"""자식이 없는 노드를 글자 순으로 돌려준다."""
parents = {parent for parent, _ in edges}
return sorted(child for _, child in edges if child not in parents)
def depths(edges: list[tuple[str, str]]) -> dict[str, int]:
"""뿌리에서 각 노드까지의 간선 수."""
children = children_of(edges)
root = find_root(edges)
result = {root: 0}
stack = [root]
while stack:
node = stack.pop()
for child in children[node]:
result[child] = result[node] + 1
stack.append(child)
return result
def height(edges: list[tuple[str, str]], node: str | None = None) -> int:
"""node(없으면 뿌리)에서 가장 먼 잎까지의 간선 수. 잎은 0이다."""
children = children_of(edges)
start = node if node is not None else find_root(edges)
def measure(current: str) -> int:
# 자식이 없으면 max의 기본값 -1에 1을 더해 0이 된다
return 1 + max((measure(child) for child in children[current]), default=-1)
return measure(start)
def subtree_size(edges: list[tuple[str, str]], node: str) -> int:
"""node와 그 자손 전체의 노드 수."""
children = children_of(edges)
return 1 + sum(subtree_size(edges, child) for child in children[node])
if __name__ == "__main__":
# 문제 페이지의 열 노드 트리
tree = [
("A", "B"), ("A", "C"), ("A", "D"),
("B", "E"), ("B", "F"), ("D", "G"),
("F", "H"), ("F", "I"), ("I", "J"),
]
assert find_root(tree) == "A"
assert leaves(tree) == ["C", "E", "G", "H", "J"]
assert depths(tree)["H"] == 3
assert height(tree) == 4
assert height(tree, "F") == 2
assert subtree_size(tree, "B") == 6
assert subtree_size(tree, "A") == len(tree) + 1
# 노드 둘짜리 트리
assert height([("R", "X")]) == 1
assert leaves([("R", "X")]) == ["X"]
print("ok")이진 탐색 트리
"""이진 탐색 트리: 찾기·넣기·지우기(자식이 둘이면 후계자로 바꿔 끼우기)와 중위 순회.
시간: 세 연산 모두 O(h) — h는 높이. 균형이 맞으면 O(log n), 한 줄이면 O(n)
공간: 노드 n개에 O(n). 재귀 깊이 O(h)
"""
from __future__ import annotations
from dataclasses import dataclass
@dataclass
class Node:
"""BST 노드: 왼쪽 부분 트리는 key보다 작고 오른쪽은 크다."""
key: int
left: Node | None = None
right: Node | None = None
def insert(root: Node | None, key: int) -> Node:
"""key를 넣은 트리의 뿌리를 돌려준다. 이미 있으면 그대로 둔다."""
if root is None:
return Node(key)
if key < root.key:
root.left = insert(root.left, key)
elif key > root.key:
root.right = insert(root.right, key)
return root
def search(root: Node | None, key: int) -> tuple[bool, int]:
"""(찾았는지, 비교 횟수)를 돌려준다."""
comparisons = 0
cursor = root
while cursor is not None:
comparisons += 1
if key == cursor.key:
return True, comparisons
cursor = cursor.left if key < cursor.key else cursor.right
return False, comparisons
def delete(root: Node | None, key: int) -> Node | None:
"""key를 지운 트리의 뿌리를 돌려준다."""
if root is None:
return None
if key < root.key:
root.left = delete(root.left, key)
return root
if key > root.key:
root.right = delete(root.right, key)
return root
if root.left is None:
return root.right
if root.right is None:
return root.left
# 자식이 둘이면 오른쪽 부분 트리의 가장 작은 값(후계자)을 복사하고 후계자를 지운다
successor = root.right
while successor.left is not None:
successor = successor.left
root.key = successor.key
root.right = delete(root.right, successor.key)
return root
def in_order(root: Node | None) -> list[int]:
"""왼쪽 → 자신 → 오른쪽 순서의 열쇠 목록(작은 순)."""
if root is None:
return []
return in_order(root.left) + [root.key] + in_order(root.right)
def height(root: Node | None) -> int:
"""간선 수로 잰 높이. 빈 트리는 -1."""
if root is None:
return -1
return 1 + max(height(root.left), height(root.right))
def build(keys: list[int]) -> Node | None:
"""빈 트리에 keys를 차례로 넣는다."""
root = None
for key in keys:
root = insert(root, key)
return root
if __name__ == "__main__":
tree = build([50, 30, 70, 20, 40, 60, 80])
tree = insert(tree, 45)
assert tree.left.right.right.key == 45
assert search(tree, 65) == (False, 3)
assert search(tree, 45) == (True, 4)
tree = build([50, 30, 70, 20, 40, 60, 80, 35, 45, 65])
tree = delete(tree, 30)
tree = delete(tree, 50)
assert tree.key == 60
assert in_order(tree) == [20, 35, 40, 45, 60, 65, 70, 80]
# 정렬된 순서로 넣으면 한 줄이 된다
assert height(build([1, 2, 3, 4, 5, 6, 7])) == 6
assert height(build([4, 2, 6, 1, 3, 5, 7])) == 2
assert delete(None, 1) is None
assert search(None, 1) == (False, 0)
print("ok")AVL·레드-블랙 트리
"""AVL 트리: 넣은 뒤 균형 인수가 ±2인 노드를 LL·LR·RR·RL 네 경우로 나눠 회전한다.
시간: 넣기 O(log n) — 높이가 약 1.44·log₂ n 이하로 묶이고, 회전은 O(1)
공간: O(n). 재귀 깊이 O(log n)
"""
from __future__ import annotations
from dataclasses import dataclass
@dataclass
class Node:
"""AVL 노드. height는 그 노드를 뿌리로 한 부분 트리의 높이(잎 1)다."""
key: int
left: Node | None = None
right: Node | None = None
height: int = 1
def height_of(node: Node | None) -> int:
"""빈 자리는 0."""
return node.height if node else 0
def update(node: Node) -> None:
"""두 자식으로 높이를 다시 잰다."""
node.height = 1 + max(height_of(node.left), height_of(node.right))
def balance_factor(node: Node) -> int:
"""왼쪽 높이 − 오른쪽 높이."""
return height_of(node.left) - height_of(node.right)
def rotate_right(node: Node) -> Node:
"""왼쪽 자식을 위로 올린다. 그 자식의 오른쪽 부분 트리는 node의 왼쪽으로 옮긴다."""
pivot = node.left
node.left = pivot.right
pivot.right = node
update(node)
update(pivot)
return pivot
def rotate_left(node: Node) -> Node:
"""rotate_right의 거울상."""
pivot = node.right
node.right = pivot.left
pivot.left = node
update(node)
update(pivot)
return pivot
def insert(node: Node | None, key: int, log: list[str]) -> Node:
"""key를 넣고 균형을 맞춘 부분 트리를 돌려준다. 회전이 일어나면 log에 경우를 적는다."""
if node is None:
return Node(key)
if key < node.key:
node.left = insert(node.left, key, log)
elif key > node.key:
node.right = insert(node.right, key, log)
else:
return node
update(node)
factor = balance_factor(node)
if factor > 1:
if key > node.left.key:
log.append(f"LR@{node.key}")
node.left = rotate_left(node.left)
else:
log.append(f"LL@{node.key}")
return rotate_right(node)
if factor < -1:
if key < node.right.key:
log.append(f"RL@{node.key}")
node.right = rotate_right(node.right)
else:
log.append(f"RR@{node.key}")
return rotate_left(node)
return node
def in_order(node: Node | None) -> list[int]:
"""작은 순 열쇠 목록."""
if node is None:
return []
return in_order(node.left) + [node.key] + in_order(node.right)
def is_balanced(node: Node | None) -> bool:
"""모든 노드의 좌우 높이 차가 1 이하인가."""
if node is None:
return True
return abs(balance_factor(node)) <= 1 and is_balanced(node.left) and is_balanced(node.right)
if __name__ == "__main__":
rotations: list[str] = []
root = None
for value in [10, 20, 30, 25, 28, 27, 5]:
root = insert(root, value, rotations)
assert rotations == ["RR@10", "LR@30", "RL@20", "LL@20"]
assert root.key == 25
assert root.height == 3
assert in_order(root) == [5, 10, 20, 25, 27, 28, 30]
assert is_balanced(root)
# 정렬된 입력 1000개도 높이가 1.44·log₂ n 아래로 묶인다
root = None
for value in range(1000):
root = insert(root, value, [])
assert is_balanced(root)
assert root.height <= 15
for keys, case in [([30, 20, 10], "LL"), ([30, 10, 20], "LR"), ([10, 20, 30], "RR"), ([10, 30, 20], "RL")]:
log: list[str] = []
root = None
for value in keys:
root = insert(root, value, log)
assert log[0].startswith(case) and root.key == 20
print("ok")B-트리
"""B-트리 넣기: 열쇠는 늘 잎에 넣고, 노드가 넘치면 가운데 열쇠를 부모로 올리며 둘로 가른다.
시간: 찾기·넣기 O(t·log_t n) — t는 노드당 열쇠 수, 디스크 읽기는 높이만큼 O(log_t n)
공간: O(n)
"""
from __future__ import annotations
from bisect import bisect_left, insort
from dataclasses import dataclass, field
@dataclass
class Node:
"""정렬된 열쇠와, 잎이 아니면 열쇠 수 + 1개의 자식."""
keys: list[int] = field(default_factory=list)
children: list[Node] = field(default_factory=list)
@property
def is_leaf(self) -> bool:
"""자식이 없으면 잎이다."""
return not self.children
class BTree:
"""노드마다 열쇠를 max_keys개까지 담는 B-트리 (max_keys=2이면 2-3 트리)."""
def __init__(self, max_keys: int = 2) -> None:
if max_keys < 2:
raise ValueError("노드에는 열쇠가 2개 이상 들어가야 한다")
self.max_keys = max_keys
self.root = Node()
self.splits = 0
def search(self, key: int) -> tuple[bool, int]:
"""(찾았는지, 읽은 노드 수)."""
node = self.root
reads = 0
while True:
reads += 1
index = bisect_left(node.keys, key)
if index < len(node.keys) and node.keys[index] == key:
return True, reads
if node.is_leaf:
return False, reads
node = node.children[index]
def insert(self, key: int) -> None:
"""key를 넣는다. 뿌리가 갈라지면 트리가 위로 한 층 자란다."""
promoted = self._insert(self.root, key)
if promoted is not None:
middle, right = promoted
self.root = Node([middle], [self.root, right])
def _insert(self, node: Node, key: int) -> tuple[int, Node] | None:
"""넣은 뒤 node가 넘쳐 갈라지면 (올릴 가운데 열쇠, 오른쪽 반쪽)을 돌려준다."""
if node.is_leaf:
insort(node.keys, key)
else:
index = bisect_left(node.keys, key)
promoted = self._insert(node.children[index], key)
if promoted is not None:
middle, right = promoted
node.keys.insert(index, middle)
node.children.insert(index + 1, right)
if len(node.keys) <= self.max_keys:
return None
return self._split(node)
def _split(self, node: Node) -> tuple[int, Node]:
"""가운데 열쇠를 떼어 내고 node를 왼쪽 반쪽으로 줄인다."""
self.splits += 1
middle_index = len(node.keys) // 2
middle = node.keys[middle_index]
right = Node(node.keys[middle_index + 1:], node.children[middle_index + 1:])
node.keys = node.keys[:middle_index]
node.children = node.children[:middle_index + 1]
return middle, right
def levels(self) -> list[list[list[int]]]:
"""층마다 노드의 열쇠 목록."""
result = []
layer = [self.root]
while layer:
result.append([node.keys for node in layer])
layer = [child for node in layer for child in node.children]
return result
if __name__ == "__main__":
tree = BTree(max_keys=2)
for value in [50, 20, 80, 10, 30, 60, 90, 40, 70]:
tree.insert(value)
assert tree.levels() == [[[50]], [[20], [80]], [[10], [30, 40], [60, 70], [90]]]
assert tree.splits == 4
tree = BTree(max_keys=2)
for value in range(10, 80, 10):
tree.insert(value)
assert tree.levels() == [[[40]], [[20], [60]], [[10], [30], [50], [70]]]
assert tree.search(55) == (False, 3)
assert tree.search(40) == (True, 1)
# 노드가 넓으면 열쇠 1만 개도 낮게 담긴다
wide = BTree(max_keys=99)
for value in range(10_000):
wide.insert(value)
assert len(wide.levels()) <= 3
assert all(wide.search(value)[0] for value in (0, 5_000, 9_999))
assert BTree().search(1) == (False, 1)
print("ok")이진 힙과 우선순위 큐
"""이진 최소 힙(우선순위 큐): 배열에 층 순서로 담고, 넣을 때 올리고(sift-up) 꺼낼 때 내린다(sift-down).
시간: 넣기·꺼내기 O(log n), 최솟값 보기 O(1), 한꺼번에 만들기(heapify) O(n)
공간: O(n) — 포인터 없이 배열 하나
"""
class MinHeap:
"""i번 칸의 자식은 2i+1·2i+2번, 부모는 (i-1)//2번이다."""
def __init__(self, values: list[int] | None = None) -> None:
self.values = list(values or [])
self.swaps = 0
# 아래층 부모부터 차례로 내리면 O(n)에 힙이 된다
for index in range(len(self.values) // 2 - 1, -1, -1):
self._sift_down(index)
def __len__(self) -> int:
return len(self.values)
def peek(self) -> int:
"""가장 작은 값을 꺼내지 않고 본다."""
if not self.values:
raise IndexError("빈 힙")
return self.values[0]
def push(self, value: int) -> None:
"""끝에 붙이고 부모보다 작은 동안 올린다."""
self.values.append(value)
index = len(self.values) - 1
while index > 0:
parent = (index - 1) // 2
if self.values[parent] <= self.values[index]:
break
self._swap(parent, index)
index = parent
def pop(self) -> int:
"""가장 작은 값을 꺼낸다. 끝 값을 뿌리로 옮겨 내린다."""
if not self.values:
raise IndexError("빈 힙")
smallest = self.values[0]
last = self.values.pop()
if self.values:
self.values[0] = last
self._sift_down(0)
return smallest
def _sift_down(self, index: int) -> None:
"""두 자식 중 작은 쪽보다 크면 자리를 바꾸며 내려간다."""
size = len(self.values)
while True:
smaller = index
for child in (2 * index + 1, 2 * index + 2):
if child < size and self.values[child] < self.values[smaller]:
smaller = child
if smaller == index:
return
self._swap(index, smaller)
index = smaller
def _swap(self, first: int, second: int) -> None:
self.values[first], self.values[second] = self.values[second], self.values[first]
self.swaps += 1
def is_min_heap(values: list[int]) -> bool:
"""모든 칸이 부모보다 작지 않은가."""
return all(values[(index - 1) // 2] <= values[index] for index in range(1, len(values)))
def top_k(stream: list[int], k: int) -> list[int]:
"""크기 k인 최소 힙으로 가장 큰 k개를 고른다. O(n log k)."""
heap = MinHeap()
for value in stream:
if len(heap) < k:
heap.push(value)
elif value > heap.peek():
# 뿌리(지금의 k등)를 내보내고 새 값을 들인다
heap.pop()
heap.push(value)
return sorted(heap.values)
if __name__ == "__main__":
heap = MinHeap()
for value in [6, 2, 8, 1, 7, 3]:
heap.push(value)
assert heap.values == [1, 2, 3, 6, 7, 8]
assert heap.pop() == 1
assert heap.pop() == 2
assert heap.values == [3, 6, 7, 8]
assert heap.swaps == 7
assert is_min_heap([1, 3, 2, 7, 4, 5])
assert not is_min_heap([2, 4, 3, 8, 1, 9])
assert is_min_heap([1, 2, 3, 4, 5, 6])
assert top_k([5, 12, 3, 8, 15, 7, 10], 3) == [10, 12, 15]
built = MinHeap([9, 4, 7, 1, 8, 2])
assert is_min_heap(built.values)
assert [built.pop() for _ in range(6)] == [1, 2, 4, 7, 8, 9]
try:
MinHeap().pop()
raise AssertionError("빈 힙에서 꺼내면 IndexError가 나야 한다")
except IndexError:
pass
print("ok")트라이
"""트라이(접두사 트리): 글자 하나가 간선 하나다. 낱말 넣기·찾기, 접두사 검색, 압축 트라이 노드 수.
시간: 길이 L인 낱말의 넣기·찾기 O(L), 접두사 p 검색 O(|p| + 결과 크기)
공간: O(전체 글자 수) — 앞부분이 같은 낱말은 노드를 함께 쓴다
"""
from __future__ import annotations
from dataclasses import dataclass, field
@dataclass
class TrieNode:
"""자식 표(글자 → 노드)와 낱말 끝 표시."""
children: dict[str, TrieNode] = field(default_factory=dict)
end: bool = False
class Trie:
"""뿌리는 빈 문자열을 뜻한다."""
def __init__(self) -> None:
self.root = TrieNode()
def insert(self, word: str) -> int:
"""word를 넣고 새로 만든 노드 수를 돌려준다."""
node = self.root
created = 0
for letter in word:
if letter not in node.children:
node.children[letter] = TrieNode()
created += 1
node = node.children[letter]
node.end = True
return created
def _walk(self, prefix: str) -> TrieNode | None:
"""prefix를 따라 내려간 노드. 길이 끊기면 None."""
node = self.root
for letter in prefix:
node = node.children.get(letter)
if node is None:
return None
return node
def contains(self, word: str) -> bool:
"""길 위에 있는 것만으로는 안 되고 끝 표시가 있어야 낱말이다."""
node = self._walk(word)
return node is not None and node.end
def starts_with(self, prefix: str) -> list[str]:
"""prefix로 시작하는 낱말을 사전 순으로."""
start = self._walk(prefix)
if start is None:
return []
words: list[str] = []
stack = [(start, prefix)]
while stack:
node, text = stack.pop()
if node.end:
words.append(text)
for letter, child in node.children.items():
stack.append((child, text + letter))
return sorted(words)
def node_count(self) -> int:
"""뿌리를 포함한 노드 수."""
return _count(self.root, compressed=False)
def compressed_node_count(self) -> int:
"""갈래도 낱말 끝도 없는 외길을 합친 압축 트라이(radix tree)의 노드 수."""
return _count(self.root, compressed=True)
def _count(node: TrieNode, compressed: bool) -> int:
"""node 아래 노드 수. 압축하면 자식 하나·끝 아님인 노드를 건너뛴다."""
total = 1
for child in node.children.values():
while compressed and len(child.children) == 1 and not child.end:
child = next(iter(child.children.values()))
total += _count(child, compressed)
return total
if __name__ == "__main__":
trie = Trie()
created = [trie.insert(word) for word in ["ten", "tea", "to", "tent", "in", "t"]]
assert created == [3, 1, 1, 1, 2, 0]
assert trie.node_count() == 9
assert trie.starts_with("te") == ["tea", "ten", "tent"]
assert trie.contains("t")
assert not trie.contains("te")
assert trie.starts_with("x") == []
sample = Trie()
for word in ["to", "tea", "ted", "ten", "in", "inn"]:
sample.insert(word)
assert sample.node_count() == 10
assert sample.contains("in") and not sample.contains("te")
radix = Trie()
for word in ["test", "team", "tea", "toast", "to"]:
radix.insert(word)
assert radix.node_count() == 11
assert radix.compressed_node_count() == 8
assert Trie().node_count() == 1
print("ok")세그먼트 트리·펜윅 트리
"""구간 합 질의: 세그먼트 트리(구간 합·값 바꾸기)와 펜윅 트리(앞부분 합·값 더하기).
시간: 두 구조 모두 질의·갱신 O(log n), 만들기 O(n) (펜윅은 여기서 O(n log n))
공간: 세그먼트 트리 O(4n), 펜윅 트리 O(n)
"""
class SegmentTree:
"""노드 1이 전체, 노드 i의 자식은 2i·2i+1이 앞·뒤 절반을 맡는다."""
def __init__(self, values: list[int]) -> None:
if not values:
raise ValueError("빈 배열")
self.size = len(values)
self.sums = [0] * (4 * self.size)
self._build(1, 0, self.size - 1, values)
def _build(self, node: int, low: int, high: int, values: list[int]) -> None:
if low == high:
self.sums[node] = values[low]
return
middle = (low + high) // 2
self._build(2 * node, low, middle, values)
self._build(2 * node + 1, middle + 1, high, values)
self.sums[node] = self.sums[2 * node] + self.sums[2 * node + 1]
def query(self, left: int, right: int) -> int:
"""values[left..right](양 끝 포함)의 합."""
return self._query(1, 0, self.size - 1, left, right)
def _query(self, node: int, low: int, high: int, left: int, right: int) -> int:
if right < low or high < left:
return 0
# 통째로 들어가는 노드는 저장된 합을 그대로 쓴다
if left <= low and high <= right:
return self.sums[node]
middle = (low + high) // 2
return self._query(2 * node, low, middle, left, right) + self._query(2 * node + 1, middle + 1, high, left, right)
def update(self, index: int, value: int) -> list[int]:
"""values[index]를 value로 바꾸고, 고친 노드 번호를 잎부터 돌려준다."""
touched: list[int] = []
self._update(1, 0, self.size - 1, index, value, touched)
return touched
def _update(self, node: int, low: int, high: int, index: int, value: int, touched: list[int]) -> None:
if low == high:
self.sums[node] = value
else:
middle = (low + high) // 2
if index <= middle:
self._update(2 * node, low, middle, index, value, touched)
else:
self._update(2 * node + 1, middle + 1, high, index, value, touched)
self.sums[node] = self.sums[2 * node] + self.sums[2 * node + 1]
touched.append(node)
class FenwickTree:
"""1부터 센다. tree[i]는 (i − lowbit(i), i] 구간의 합이다."""
def __init__(self, values: list[int]) -> None:
self.tree = [0] * (len(values) + 1)
for position, value in enumerate(values, start=1):
self.add(position, value)
def add(self, position: int, delta: int) -> list[int]:
"""position의 값에 delta를 더하고, 고친 칸을 돌려준다."""
touched = []
while position < len(self.tree):
self.tree[position] += delta
touched.append(position)
position += position & -position
return touched
def prefix_sum(self, count: int) -> int:
"""앞 count칸의 합."""
total = 0
while count > 0:
total += self.tree[count]
count -= count & -count
return total
def range_sum(self, left: int, right: int) -> int:
"""left..right칸(1부터, 양 끝 포함)의 합 = 두 앞부분 합의 차."""
return self.prefix_sum(right) - self.prefix_sum(left - 1)
if __name__ == "__main__":
values = [3, 1, 4, 1, 5, 9, 2, 6]
segment = SegmentTree(values)
assert segment.sums[1] == 31
assert segment.query(1, 5) == 20
assert segment.query(0, 7) == sum(values)
assert segment.update(5, 4) == [13, 6, 3, 1]
assert segment.sums[1] == 26
assert segment.query(4, 7) == 17
fenwick = FenwickTree(values)
assert fenwick.tree[1:] == [3, 4, 4, 9, 5, 14, 2, 31]
assert fenwick.prefix_sum(6) == 23
assert fenwick.range_sum(3, 7) == 21
assert fenwick.add(3, 2) == [3, 4, 8]
assert fenwick.range_sum(3, 3) == 6
# 모든 구간을 단순 합과 비교한다
fresh = SegmentTree(values)
fresh_fenwick = FenwickTree(values)
for left in range(len(values)):
for right in range(left, len(values)):
expected = sum(values[left:right + 1])
assert fresh.query(left, right) == expected
assert fresh_fenwick.range_sum(left + 1, right + 1) == expected
assert SegmentTree([7]).query(0, 0) == 7
print("ok")