[snippet] segment-tree.py

백준 2042번. 구간 합 구하기에 대한 풀이를 Segment Tree의 스니펫 느낌으로 작성.

1. l, r, node 해석법 : segment_tree[node]에 arr[l] ~ arr[r]의 구간합이 들어있다.

• init(0, N-1, 1) : segment_tree[1]에 arr[0] ~ arr[N-1] 전체의 부분합을 담겠다.
• 헷갈리기 쉬운 것이, arr은 0-indexed이지만, segment_tree는 1-indexed라는 점이다.

2. segment_tree 읽는 법

• arr : [1, 2, 3, 4, 5] 일 때,
• segment_tree : [0, 15, 6, 9, 3, 3, 4, 5, 1, 2, 0, 0, 0, 0, 0, 0]
• 1부터 시작해서, node*2 자리에 좌측 부분합이 있고, 그 옆(node*2+1)에 우측 부분합이 있다고 읽으면 수월하다.

import sys
from math import ceil, log
input = lambda: sys.stdin.readline().rstrip()

def mid(l, r):
    return (l+r) // 2

###############################################

def init(l, r, node):
    if l == r:
        segment_tree[node] = arr[l]
        return
    init(l, mid(l,r), node*2)
    init(mid(l,r)+1, r, node*2+1)
    segment_tree[node] = segment_tree[node*2] + segment_tree[node*2+1]

def update(l, r, node, IDX, DIFF):
    if not (l <= IDX <= r):
        return
    segment_tree[node] += DIFF
    if l == r:
        return
    update(l, mid(l,r), node*2, IDX, DIFF)
    update(mid(l,r)+1, r, node*2+1, IDX, DIFF)

def interval_sum(l, r, node, LEFT, RIGHT):
    if r < LEFT or RIGHT < l: # [l, r]이 [LEFT, RIGHT]를 완전히 벗어남. 
        return 0
    if LEFT <= l and r <= RIGHT: # [l, r]이 [LEFT, RIGHT] 안에 완전히 포함됨. 
        return segment_tree[node]
    lsum = interval_sum(l, mid(l,r), node*2, LEFT, RIGHT)
    rsum = interval_sum(mid(l,r)+1, r, node*2+1, LEFT, RIGHT)
    return lsum + rsum

N, M, K = map(int, input().split())
arr = []
for _ in range(N):
    arr.append(int(input()))

segment_tree = [0] * (2**ceil(log(N, 2)+1)) # <= 4N
init(0, N-1, 1) # 1-indexed로 해야, 2*1, 2*1+1이 됨. 

###############################################

for _ in range(M+K):
    a, b, c = map(int, input().split())
    if a == 1:
        b -= 1
        update(0, N-1, 1, b, c-arr[b])
        arr[b] = c 
    else:
        b -= 1; c -= 1 # IDX, LEFT, RIGHT는 arr 인덱스 기준이므로, 0-indexed
        print(interval_sum(0, N-1, 1, b, c))