(BOJ10999 구간 합 구하기 2
문제 설명
- 처음 풀어본 lazy propagation 세그먼트 트리 문제
- lazy값을 잘 활용하는 것이 이 문제 해결의 핵심인 것 같다.
문제 풀이
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()
메서드를 작성 및 활용했다.