연습 문제

트리와 힙 문제는 코딩 테스트와 면접에서 자주 나옵니다. 아래 네 문제는 순회, BST 성질, 크기를 제한한 힙, 여러 줄을 합치는 힙이라는 대표 유형을 하나씩 다룹니다. 먼저 풀이 방향을 스스로 생각해 본 뒤 접근법과 코드를 확인하세요. 모든 코드는 아래 Node 정의를 공유합니다.

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

문제 1. 층마다 가장 큰 값

이진 트리가 주어질 때, 루트부터 깊이 순서로 각 깊이에서 가장 큰 키를 리스트로 돌려주세요. 빈 트리면 빈 리스트입니다. 예를 들어 루트 4, 그 자식 9와 2, 손자 3, 5, 7이면 [4, 9, 7]입니다.

접근: 깊이별로 묶어야 하므로 레벨 순회가 자연스럽습니다. 큐에서 한 번에 "현재 큐 길이"만큼 꺼내면 정확히 한 층을 처리하게 됩니다. 시간 O(n), 공간은 가장 넓은 층의 크기입니다.

from collections import deque


def level_maxima(root):
    if root is None:
        return []
    result, queue = [], deque([root])
    while queue:
        best = None
        for _ in range(len(queue)):        # 이번 층만큼만 꺼낸다
            node = queue.popleft()
            best = node.key if best is None else max(best, node.key)
            queue.extend(c for c in (node.left, node.right) if c)
        result.append(best)
    return result


tree = Node(4, Node(9, Node(3), Node(5)), Node(2, None, Node(7)))
print(level_maxima(tree))  # [4, 9, 7]

문제 2. BST 구간 합

이진 탐색 트리와 두 정수 lo, hi가 주어집니다. 키가 lo 이상 hi 이하인 노드들의 키 합을 구하세요. 단, 구간 밖임이 확실한 서브트리는 방문하지 않아야 합니다.

접근: BST 성질을 이용해 가지치기를 합니다. 현재 키가 lo보다 작으면 왼쪽 서브트리는 모두 더 작으므로 건너뛰고, hi보다 크면 오른쪽을 건너뜁니다. 방문하는 노드는 O(h + 결과 노드 수) 정도입니다. 깊은 트리를 대비해 스택으로 씁니다.

def range_sum(root, lo, hi):
    total, stack = 0, [root]
    while stack:
        node = stack.pop()
        if node is None:
            continue
        if lo <= node.key <= hi:
            total += node.key
        if node.key > lo:          # 왼쪽에 lo 이상인 값이 있을 수 있음
            stack.append(node.left)
        if node.key < hi:          # 오른쪽에 hi 이하인 값이 있을 수 있음
            stack.append(node.right)
    return total


bst = Node(10, Node(5, Node(3), Node(7)), Node(15, None, Node(18)))
print(range_sum(bst, 6, 15))  # 7 + 10 + 15 = 32

문제 3. 실시간 K번째 높은 점수

게임 서버에 점수가 하나씩 들어옵니다. 점수가 들어올 때마다 "지금까지 들어온 점수 중 K번째로 높은 점수"를 출력하세요. 아직 K개가 안 되었으면 None을 출력합니다.

접근: 크기를 K로 제한한 최소 힙을 유지합니다. 힙에는 상위 K개만 남기고, 루트가 바로 K번째로 높은 점수입니다. 새 점수가 루트보다 크면 루트와 바꿉니다(heapq.heappushpop이나 heapreplace). 점수당 O(log K), 메모리 O(K)로 전체를 정렬하는 것보다 훨씬 가볍습니다.

import heapq


def kth_highest_stream(scores, k):
    heap, out = [], []
    for s in scores:
        if len(heap) < k:
            heapq.heappush(heap, s)
        elif s > heap[0]:
            heapq.heapreplace(heap, s)   # 루트를 빼고 s 를 넣음
        out.append(heap[0] if len(heap) == k else None)
    return out


print(kth_highest_stream([40, 10, 70, 20, 90, 50], 3))
# [None, None, 10, 20, 40, 50]

문제 4. 여러 서버의 로그 합치기

서버 k대가 각자 시간순으로 정렬된 로그 (시각, 메시지) 리스트를 갖고 있습니다. 전체를 시간순 하나의 리스트로 합치세요. 시각이 같으면 서버 번호가 작은 쪽을 먼저 둡니다.

접근: 각 리스트의 "다음 후보" 하나씩만 힙에 넣고, 가장 이른 것을 꺼낼 때마다 그 리스트의 다음 원소를 넣습니다. 힙 원소를 (시각, 서버 번호, 위치)로 두면 같은 시각일 때 서버 번호로 순서가 정해지고, 메시지끼리 비교하는 일이 없습니다. 전체 원소가 N개면 O(N log k)입니다. 표준 라이브러리의 heapq.merge도 같은 방식입니다.

import heapq


def merge_logs(logs):
    heap = [(lst[0][0], i, 0) for i, lst in enumerate(logs) if lst]
    heapq.heapify(heap)
    merged = []
    while heap:
        t, i, j = heapq.heappop(heap)
        merged.append(logs[i][j])
        if j + 1 < len(logs[i]):
            heapq.heappush(heap, (logs[i][j + 1][0], i, j + 1))
    return merged


logs = [
    [(1, "a:start"), (5, "a:done")],
    [(2, "b:start"), (5, "b:retry"), (9, "b:done")],
    [],
    [(3, "d:start")],
]
print([m for _, m in merge_logs(logs)])
# ['a:start', 'b:start', 'd:start', 'a:done', 'b:retry', 'b:done']

풀이를 점검하는 습관

• 빈 트리, 노드 하나, 한쪽으로만 긴 트리를 꼭 시험합니다.
• 힙에 튜플을 넣을 때 앞 원소가 같으면 다음 원소를 비교하므로, 비교할 수 없는 값이 뒤에 오지 않게 합니다.
• 크기 K 힙 문제에서 "K번째로 큰 값"에는 최소 힙, "K번째로 작은 값"에는 최대 힙(값에 -1을 곱함)을 씁니다.

정리

층 단위 문제는 레벨 순회, 범위 문제는 BST 가지치기, "상위 K개"는 크기 K의 최소 힙, 여러 정렬된 입력을 합치는 문제는 각 입력의 선두를 담은 힙으로 풉니다. 이 네 가지 틀을 익히면 비슷한 문제 대부분에 바로 적용할 수 있습니다.