[BOJ] 2042 - 구간 합 구하기 (Python)

 

 

아이디어

 

세그먼트 트리로 구간 합을 저장해두고 변경이나 계산이 필요한 구간에 대해 빠르게 처리한다.

 

 

풀이

 

구간 합을 계산하기 위해 N보다 크면서 가장 작은 2의 거듭제곱의 두 배 크기로 0으로 채워진 트리 리스트 tree를 만든다.

트리의 리프 노드들에는 순서대로 N개의 수를 입력 받고, for문에서 각 부모 노드에 두 자식 노드의 합을 저장한다.

이후 for문에서는 a가 1인 경우 해당 트리 리프 노드의 수를 변경한 후 함수 update를 호출하고,

a가 2인 경우 구간 합을 구하는 함수 sum_up을 호출한다.

 

함수 update에서는 리프 노드의 부모 노드부터 루트 노드까지 다시 계산하고, 재귀 호출로 루트 노드를 벗어나면 종료한다.

 

함수 sum_up에서는 리프 노드에서부터 차례로 해당 노드의 부모 노드로 올라가며 구간 합을 결과값 res에 더해 나간다.

부모 노드의 값은 두 자식 노드 값의 합이므로 s%2 == 1 또는 e%2 == 0처럼 형제 노드 중 한 노드만 포함되어 있는 경우

해당 노드 값을 res에 더한 후 인덱스를 이동하여 부모 노드의 값을 더하지 않도록 한다.

 

import sys

input = sys.stdin.readline

def update(idx):
    if idx < 1:
        return
    tree[idx] = tree[idx*2]+tree[idx*2+1]
    update(idx//2)

def sum_up(s, e):
    res = 0
    while s <= e:
        if s%2:
            res += tree[s]
            s += 1
        if e%2 == 0:
            res += tree[e]
            e -= 1
        s, e = s//2, e//2
    return res

N, M, K = map(int, input().split())
n = 1
while n < N:
    n *= 2
tree = [0]*(n*2)
tree[n:n+N] = [int(input()) for _ in range(N)]
for i in range(n-1, 0, -1):
    tree[i] = tree[i*2]+tree[i*2+1]
for _ in range(M+K):
    a, b, c = map(int, input().split())
    if a == 1:
        tree[n+b-1] = c
        update((n+b-1)//2)
    else:
        print(sum_up(b+n-1, c+n-1))