[BOJ] 2268 - 수들의 합 7 (Python)

 

 

아이디어

 

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

 

 

풀이

 

1275번 커피숍2 문제 풀이와 유사하다.

 

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

for문에서 a가 0이 아닌 경우(1인 경우) 해당 트리 리프 노드의 수를 변경한 후 함수 modify를 호출하고,

a가 0인 경우 구간 합을 구하는 함수 sum_up을 호출하는데, b가 c보다 더 크다면 b와 c를 바꿔 c에 더 큰 값이 오도록 한다.

 

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

 

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

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

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

 

import sys

input = sys.stdin.readline

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

def modify(idx):
    while idx > 0:
        tree[idx] = tree[idx*2]+tree[idx*2+1]
        idx //= 2

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