연습 문제

이 장의 문제는 구간 트리와 펜윅 트리의 대표 유형을 연습하도록 새로 만든 것입니다. 먼저 혼자 풀어 본 뒤 접근 방법과 풀이를 확인하세요. 모든 풀이는 표준 라이브러리만 쓰는 Python 코드입니다.

문제 1. 일별 판매량 장부

어떤 매장의 n일치 판매량이 배열로 주어집니다. 두 종류의 요청을 차례로 처리합니다. ("set", d, v)는 d일의 판매량을 v로 고치고, ("sum", l, r)은 l일부터 r-1일까지의 판매량 합을 출력합니다. n과 요청 수는 각각 최대 200,000입니다.

접근: 질의가 구간 합이므로 펜윅 트리가 가장 간단합니다. 펜윅 트리는 더하기만 받으므로 원래 배열을 따로 들고 있다가 차이를 더합니다. 처음 만들 때 add를 n번 부르면 O(n log n)이지만, 아래처럼 자기 값을 바로 위 부모에게 넘기면 O(n)에 만들 수 있습니다.

def solve(sales, requests):
    n = len(sales)
    tree = [0] + sales[:]
    for i in range(1, n + 1):              # O(n) 만들기
        parent = i + (i & -i)
        if parent <= n:
            tree[parent] += tree[i]

    def add(i, delta):
        i += 1
        while i <= n:
            tree[i] += delta
            i += i & -i

    def prefix(count):
        total = 0
        while count > 0:
            total += tree[count]
            count -= count & -count
        return total

    out = []
    for kind, x, y in requests:
        if kind == "set":
            add(x, y - sales[x])
            sales[x] = y
        else:
            out.append(prefix(y) - prefix(x))
    return out


print(solve([4, 0, 7, 2, 5], [("sum", 1, 4), ("set", 1, 6), ("sum", 0, 5)]))   # [9, 24]

문제 2. 가장 추운 구간

센서 n개의 온도가 주어집니다. ("set", i, t)는 센서 i의 온도를 t로 바꾸고, ("min", l, r)은 센서 l부터 r-1까지 가운데 가장 낮은 온도를 답합니다.

접근: 최솟값은 빼기로 되돌릴 수 없으므로 펜윅 트리로는 풀기 어렵습니다. 합치는 연산을 min, 항등원을 무한대로 둔 구간 트리를 씁니다.

def coldest(temps, requests):
    n = len(temps)
    INF = float("inf")
    tree = [INF] * n + temps
    for i in range(n - 1, 0, -1):
        tree[i] = min(tree[2 * i], tree[2 * i + 1])

    out = []
    for kind, x, y in requests:
        if kind == "set":
            i = x + n
            tree[i] = y
            while i > 1:
                i //= 2
                tree[i] = min(tree[2 * i], tree[2 * i + 1])
        else:
            lo, hi, best = x + n, y + n, INF
            while lo < hi:
                if lo & 1:
                    best = min(best, tree[lo]); lo += 1
                if hi & 1:
                    hi -= 1; best = min(best, tree[hi])
                lo //= 2; hi //= 2
            out.append(best)
    return out


print(coldest([3, -2, 5, 0, 1], [("min", 2, 5), ("set", 3, 4), ("min", 2, 5), ("min", 0, 5)]))
# [0, 1, -2]

문제 3. 순서가 뒤바뀐 쌍 세기

정수 배열에서 i < j이면서 a[i] > a[j]인 쌍(역순 쌍)의 개수를 구합니다. 길이는 최대 200,000이고 값의 범위는 매우 넓습니다.

접근: 왼쪽에서 오른쪽으로 훑으면서, 지금까지 본 값 가운데 현재 값보다 큰 값이 몇 개인지 세어 더합니다. 값의 범위가 넓으므로 먼저 정렬된 순위로 바꾸는 좌표 압축을 하고, 순위별 개수를 펜윅 트리에 담습니다. 전체 O(n log n)입니다.

def count_inversions(a):
    ranks = {v: i + 1 for i, v in enumerate(sorted(set(a)))}
    m = len(ranks)
    tree = [0] * (m + 1)
    seen = inversions = 0
    for v in a:
        r = ranks[v]
        smaller_or_equal, i = 0, r
        while i > 0:
            smaller_or_equal += tree[i]
            i -= i & -i
        inversions += seen - smaller_or_equal
        while r <= m:
            tree[r] += 1
            r += r & -r
        seen += 1
    return inversions


print(count_inversions([8, 4, 2, 1]))         # 6
print(count_inversions([3, 1, 2, 3, 1]))      # 5

문제 4. 구간 할인과 매출 합계

상품 n개의 가격이 있습니다. ("add", l, r, v)는 l번부터 r-1번 상품의 가격에 모두 v를 더하고(할인이면 음수), ("sum", l, r)은 그 구간의 가격 합을 답합니다.

접근: 구간 갱신과 구간 질의가 모두 있으므로 지연 전파 구간 트리를 씁니다. 구간을 완전히 덮는 노드에서는 합에 v * 길이를 더하고 lazy에 v를 쌓아 둔 뒤 멈춥니다. 자식으로 내려가야 할 때만 push로 lazy를 넘깁니다.

def solve_lazy(prices, requests):
    n = len(prices)
    tree, lazy = [0] * (4 * n), [0] * (4 * n)

    def build(node, lo, hi):
        if hi - lo == 1:
            tree[node] = prices[lo]
            return
        mid = (lo + hi) // 2
        build(2 * node, lo, mid); build(2 * node + 1, mid, hi)
        tree[node] = tree[2 * node] + tree[2 * node + 1]

    def apply(node, lo, hi, v):
        tree[node] += v * (hi - lo)
        lazy[node] += v

    def push(node, lo, hi):
        if lazy[node]:
            mid = (lo + hi) // 2
            apply(2 * node, lo, mid, lazy[node])
            apply(2 * node + 1, mid, hi, lazy[node])
            lazy[node] = 0

    def add(node, lo, hi, l, r, v):
        if r <= lo or hi <= l:
            return
        if l <= lo and hi <= r:
            apply(node, lo, hi, v)
            return
        push(node, lo, hi)
        mid = (lo + hi) // 2
        add(2 * node, lo, mid, l, r, v); add(2 * node + 1, mid, hi, l, r, v)
        tree[node] = tree[2 * node] + tree[2 * node + 1]

    def total(node, lo, hi, l, r):
        if r <= lo or hi <= l:
            return 0
        if l <= lo and hi <= r:
            return tree[node]
        push(node, lo, hi)
        mid = (lo + hi) // 2
        return total(2 * node, lo, mid, l, r) + total(2 * node + 1, mid, hi, l, r)

    build(1, 0, n)
    out = []
    for req in requests:
        if req[0] == "add":
            add(1, 0, n, req[1], req[2], req[3])
        else:
            out.append(total(1, 0, n, req[1], req[2]))
    return out


print(solve_lazy([10, 20, 30, 40, 50], [("add", 1, 4, -5), ("sum", 0, 3), ("sum", 2, 5)]))
# [50, 110]

정리

구간 합에 점 갱신이면 펜윅 트리, 최솟값 · 최댓값이면 구간 트리, 구간 갱신까지 섞이면 지연 전파 구간 트리를 먼저 떠올립니다. "앞에서 본 값 중 몇 개가 더 큰가"처럼 개수를 세는 문제는 좌표 압축과 펜윅 트리의 조합으로 자주 풀립니다.