[백준] 2042번 구간 합 구하기(Python)

백준 2042 구간 합 구하기

세그먼트 트리에 관한 글

세그먼트 트리에 대해서 더 깊게 공부할 필요를 느낀 문제다. 구간합 문제를 어떻게 접근해야 할지 정형화된 공부가 필요했다. 구간 합 문제를 선형으로 접근했을 때와 트리로 접근했을 때의 시간 복잡도는 각각 O(n)과 O(logn)으로 세그먼트 트리가 선형으로 접근했을 때보다 더 빠르다. 단, 구간 합을 저장하는 리스트가 있는 경우에는 리스트를 만들 때는 O(n)이고 구간 합을 구할 때는 O(1)이지만, 수정할 때 O(n)의 시간 복잡도를 갖는다.

1. 일반적인 선형 접근법: 구간 합에 접근할 때 O(n)으로 느림
2. 합을 저장해두는 리스트: 리스트를 제작, 수정할 땐, O(n) 하지만 접근할 땐 O(1)
3. Segment Tree: 수정 , 구간 합 등 모두 O(logn)의 시간 복잡도를 필요로 함

import sys

def init(start, end, nodeNum):
    if start == end:
        tree[nodeNum] = inputList[start]
        return tree[nodeNum]
    mid = (start + end) // 2
    tree[nodeNum] = init(start, mid, nodeNum * 2) + init(mid + 1, end, nodeNum * 2 + 1)
    return tree[nodeNum]

def tree_sum(start, end, nodeNum, left, right):
    if left > end or right < start:
        return 0
    if left <= start and end <= right:
        return tree[nodeNum]
    mid = (start + end) // 2
    return tree_sum(start, mid, nodeNum * 2, left, right) + tree_sum(mid + 1, end, nodeNum * 2 + 1, left, right)

def update(start, end, nodeNum, idx, fixed):
    if idx < start or idx > end:
        return
    tree[nodeNum] += fixed
    if start == end:
        return
    mid = (start + end) // 2
    update(start, mid, nodeNum * 2, idx, fixed)
    update(mid + 1, end, nodeNum * 2 + 1, idx, fixed)

input = sys.stdin.readline
n, m, k = map(int, input().split())
tree = [0 for _ in range(4 * n)]
inputList = [0]

for _ in range(n):
    inputList.append(int(input().strip()))
init(1, n, 1)
for _ in range(k + m):
    a, b, c = map(int, input().split())
    if a == 1:
        tmp = c - inputList[b]
        inputList[b] = c
        update(1, n, 1, b, tmp) 
    else:
        print(tree_sum(1, n, 1, b, c))