[Python] 백준 2042번 - 구간 합 구하기

📖 문제

2042번: 구간 합 구하기

 

🔎 풀이

세그먼트 트리를 활용하는 대표적인 문제입니다. 세그먼트를 처음 접하는 분들은 아래 블로그에 개념과 코드 설명이 잘되어 있으니 참고 부탁드립니다.

[자료구조] 세그먼트 트리 (Segment Tree)

 

위 블로그에서 개념과 코드를 보고 오셨다면, 해당 문제는 어렵지 않게 풀 수 있습니다. 세그먼트 트리를 만들어주는 함수, 세그먼트 트리의 원소를 업데이트 하는 함수, 구간을 합하는 함수 3가지로 구성하면 됩니다.

 

아래 코드에서 주의할 점은 트리를 업데이트한 후에, 기존 배열도 업데이트 해주어야 한다는 것입니다.

 

💻 코드

import sys
input=sys.stdin.readline


# start, end: 숫자 배열의 인덱스
# idx: 세그먼트 트리의 인덱스 (1부터)
def make_tree(start, end, idx):
    if start == end:
        tree[idx] = arr[start]
        return tree[idx]

    mid = (start + end) // 2
    tree[idx] = make_tree(start, mid, idx * 2) + make_tree(mid + 1, end, idx * 2 + 1)

    return tree[idx]


# target: 수정할 값의 숫자 배열 인덱스
# value: 수정할 값 (기존 값에 얼만큼 더해야 하는지의 값)
def update_tree(start, end, idx, target, value):
    if target < start or target > end:
        return
    
    tree[idx] += value
    if start == end:
        return
    
    mid = (start + end) // 2
    update_tree(start, mid, idx * 2, target, value)
    update_tree(mid + 1, end, idx * 2 + 1, target, value)


# left, right: 구하고자 하는 범위
def sum_tree(start, end, idx, left, right):
    if right < start or left > end:
        return 0

    if left <= start and right >= end:
        return tree[idx]

    mid = (start + end) // 2
    return sum_tree(start, mid, idx * 2, left, right) + sum_tree(mid + 1, end, idx * 2 + 1, left, right)


N, M, K = map(int, input().split())
arr = []
tree = [0] * (N * 4)  # 충분한 트리 공간 확보
for _ in range(N):
    arr.append(int(input()))

make_tree(0, N-1, 1)
for _ in range(M + K):
    a, b, c = map(int, input().split())
    if a == 1:
        update_tree(0, N-1, 1, b-1, c - arr[b-1])
        arr[b-1] = c  # 기존 숫자 배열 값 변경
    else:
        print(sum_tree(0, N-1, 1, b-1, c-1))