[BOJ] 1275 - 커피숍2 (Python)

 

 

아이디어

 

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

 

 

풀이

 

2042번 구간 합 구하기 문제와 유사하다.

 

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

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

이후 for문에서 x~y의 구간 합을 구하는 함수 sum_up을 호출하여 그 반환값을 출력하고,

a의 위치에 해당하는 트리 리프 노드의 수를 변경한 후 함수 update를 호출한다.

이때 y에 더 큰 값이 오도록 x가 y보다 더 크다면 x와 y의 값을 서로 바꾼다.

 

함수 update에서는 while문을 반복하며 리프 노드의 부모 노드부터 두 자식 노드의 값을 더해 다시 계산한다.

 

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

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

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

 

import sys

input = sys.stdin.readline

def update(idx):
    while idx > 0:
        tree[idx] = tree[idx*2]+tree[idx*2+1]
        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, Q = map(int, input().split())
n = 1
while n < N:
    n *= 2
tree = [0]*(n*2)
tree[n:n+N] = list(map(int, input().split()))
for i in range(n-1, 0, -1):
    tree[i] = tree[i*2]+tree[i*2+1]
for _ in range(Q):
    x, y, a, b = map(int, input().split())
    if x > y:
        x, y = y, x
    print(sum_up(x+n-1, y+n-1))
    tree[a+n-1] = b
    update((a+n-1)//2)