(BOJ10999 구간 합 구하기 2

Link copied to clipboard

문제 설명

Link copied to clipboard
  • 처음 풀어본 lazy propagation 세그먼트 트리 문제
    • lazy값을 잘 활용하는 것이 이 문제 해결의 핵심인 것 같다.

문제 풀이

Link copied to clipboard
class SegmentTree:
    def __init__(self, N: int, seq: list[int]) -> None:
        self.N = N
        self.seq = [0] + seq
        self.tree = [0] * (1 << (ceil(log2(N)) + 1))
        self.lazy = [0] * (1 << (ceil(log2(N)) + 1))
        self.init_tree(1, 1, N)

    def mid(self, start: int, end: int) -> int:
        return (start + end) // 2

    def init_tree(self, node: int, start: int, end: int) -> None:
        if start >= end:
            self.tree[node] = self.seq[start]
            return
        mid = self.mid(start, end)
        self.init_tree(node * 2, start, mid)
        self.init_tree(node * 2 + 1, mid + 1, end)
        self.tree[node] = self.tree[node * 2] + self.tree[node * 2 + 1]

    def update_lazy(self, node: int, start: int, end: int) -> None:
        if self.lazy[node]:
            self.tree[node] += (end - start + 1) * self.lazy[node]
            if start != end:
                self.lazy[node * 2] += self.lazy[node]
                self.lazy[node * 2 + 1] += self.lazy[node]
            self.lazy[node] = 0

    def update_range(
        self, node: int, start: int, end: int, left: int, right: int, diff: int
    ) -> None:
        self.update_lazy(node, start, end)
        if right < start or end < left:
            return
        if left <= start and end <= right:
            self.tree[node] += (end - start + 1) * diff
            if start != end:
                self.lazy[node * 2] += diff
                self.lazy[node * 2 + 1] += diff
            return
        mid = self.mid(start, end)
        self.update_range(node * 2, start, mid, left, right, diff)
        self.update_range(node * 2 + 1, mid + 1, end, left, right, diff)
        self.tree[node] = self.tree[node * 2] + self.tree[node * 2 + 1]

    def query(self, node: int, start: int, end: int, left: int, right: int) -> int:
        self.update_lazy(node, start, end)
        if right < start or end < left:
            return 0
        if left <= start and end <= right:
            return self.tree[node]
        mid = self.mid(start, end)
        left_sum = self.query(node * 2, start, mid, left, right)
        right_sum = self.query(node * 2 + 1, mid + 1, end, left, right)
        return left_sum + right_sum
  • 우선 문제 풀이 코드부터 보면 다음과 같다.
  • 처음 세그먼트 트리를 만들기 위한 init_tree()메서드를 선언 및 실행했고,
  • 주어지는 입력을 위하여
    • update_range(), query()메서드를 기반으로 동작한다.
    • 나머지 부분은 일반적인 segment tree와 비슷하게 동작하는데,
    • lazy propagation을 위한 update_lazy()메서드를 작성 및 활용했다.