백준 2042 - 구간 합 구하기 (Python)

이번 문제에서 세그먼트 트리를 사용하는 이유는 구간합을 빠르게 검색하고 수정하기 위해서입니다.
개념 자체는 쉬우니 바로 코드를 보고 이해하시면 됩니다.

전체코드

import sys

input = sys.stdin.readline

N, M, K = map(int,input().split())
leaf = [int(input()) for _ in range(N)]
segment_tree = [0] * (N * 4)

def make_segment_tree(s, e, index):
    if s == e:
        segment_tree[index] = leaf[s-1]
        return segment_tree[index]
	
    mid = (s + e) // 2
    segment_tree[index] = make_segment_tree(s, mid, index*2) + make_segment_tree(mid+1, e, index*2+1)
    
    return segment_tree[index]

def find(s, e, i, l, r):
    if s > r or e < l:
        return 0
        
    if s >= l and e <= r:
        return segment_tree[i]

    mid = (s + e) // 2
    
    return find(s, mid, i * 2, l, r) + find(mid + 1, e, i * 2 + 1, l, r)

def update(s, e, i, update_idx, update_data):
    if s > update_idx or e < update_idx:
        return
    
    segment_tree[i] += update_data
	
    if s == e:
        return

    mid = (s + e) // 2
    update(s, mid, i * 2, update_idx, update_data)
    update(mid+1, e, i * 2 + 1, update_idx, update_data)

make_segment_tree(1, N, 1)

for _ in range(M + K):
    a, b, c = map(int, input().split())
    
    if a == 1:
        temp = c - leaf[b - 1]
        leaf[b - 1] = c
        update(1, N, 1, b, temp)
    elif a == 2:
        print(find(1, N, 1, b, c))