[snippet] fenwick-tree.py

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

1. fenwick_tree 읽는 법

• arr : [0, 1, 2, 3, 4, 5] (1-indexed) 일 때,
• fenwick_tree : [0, 1, 3, 3, 10, 5]
• 홀수 인덱스에는 실제 리프 노드 값이 들어간다.
• 짝수 인덱스에는 처음부터 해당 인덱스까지의 부분 합이 들어간다. (fenwick_tree[4]에는 arr[1~4]의 합이 들어있다.)

2. segment_tree 스니펫과 fenwick_tree 스니펫의 차이점

• fenwick_tree에는 init()이 없다.
• segment_tree에서는 arr 기준으로 인덱스가 들어갔지만, fenwick_tree의 update, interval_sum 함수에서는 모두 fenwick_tree를 기준으로 한다.
  • 이번 문제에서는 arr까지 1-indexed로 맞춰줘서 잘 넘어갔지만, arr이 0-indexed라면 헷갈릴 수 있다.

import sys
input = lambda: sys.stdin.readline().rstrip()

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

def update(i, diff):
    while i <= N:
        fenwick_tree[i] += diff
        i += (i & -i)

def interval_sum(start, end):
    def prefix_sum(i): # [1~i]의 합
        ret = 0
        while i > 0:
            ret += fenwick_tree[i]
            i -= (i & -i)
        return ret
    return prefix_sum(end) - prefix_sum(start-1)

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

N, M, K = map(int, input().split())
fenwick_tree = [0] * (N+1) # 1-indexed 여야 함. 0 & 0을 할 수는 없으니까.

arr = [0] * (N+1)
for i in range(1, N+1):
    x = int(input())
    arr[i] = x
    update(i, x)
    
for i in range(M+K):
    a, b, c = map(int, input().split())
    if a == 1:
        update(b, c-arr[b])
        arr[b] = c
    else:
        print(interval_sum(b, c))