[BOJ 2042] 구간 합 구하기(Python)

---

문제

구간 합 구하기

---

문제 해설

세그먼트 트리의 대표적인 문제입니다. 자세한 설명은 여기

---

풀이 코드

# 6:24
import sys

input = sys.stdin.readline

# 세그먼트 트리 초기화
def init(node, start, end):
    if start == end:
        tree[node] = arr[start]
        return tree[node]
    mid = (start + end) // 2
    # 리프 노드가 아닐 때 자식 노드의 리턴 값 저장
    tree[node] += init(node * 2, start, mid) + init(node * 2 + 1, mid + 1, end)
    return tree[node]


# 구간 합
def prepix_sum(start, end, left, right, node):
    # 필요 없는 구간이므로 버림.
    if start > right or end < left:
        return 0
    # 필요한 구간은 해당 tree 값 리턴
    # 굳이 리프까지 갈 필요 X
    if left <= start and end <= right:
        return tree[node]
    mid = (start + end) // 2
    # init과 마찬가지로 자식 노드의 리턴 값을 더해 리턴
    return prepix_sum(start, mid, left, right, node * 2) + prepix_sum(
        mid + 1, end, left, right, node * 2 + 1
    )


# 값 변경하기, 반드시 한 줄로 탐색을 한다.
# 해당 인덱스를 리프 노드에서 찾은 뒤, 리턴되며 값을 전부 변경
# 리턴 값은 원래의 값과의 차이
def update(index, start, end, node, value):
    # 리프 노드를 찾았고, 값의 차이를 리턴
    if index == start == end:
        dif = value - tree[node]
        tree[node] = value
        return dif
    mid = (start + end) // 2
    dif = 0
    if index <= mid:
        dif = update(index, start, mid, node * 2, value)
        tree[node] += dif
    else:
        dif = update(index, mid + 1, end, node * 2 + 1, value)
        tree[node] += dif
    return dif


arr = []
tree = [0] * 3000000  # 구간 합 저장.
n, m, k = map(int, input().split())
# 초기 값 순서대로 저장
for _ in range(n):
    arr.append(int(input()))
init(1, 0, n - 1)

for _ in range(m + k):
    a, b, c = map(int, input().split())
    if a == 1:  # 변경
        update(b - 1, 0, n - 1, 1, c)
    else:  # 출력
        print(prepix_sum(0, n - 1, b - 1, c - 1, 1))

---