[백준 2042] 구간 합 구하기

풀이

세그먼트 트리를 이용한다

 

구글에서 검색해서 나오는 세그먼트 트리 방법들은 arr에 index를 이용하여 구현하는데 나는 그게 싫어서 class안에다가 range를 넣었다. 그러고 나니 생각보다 코드가 깔끔하게 나왔다.

 

import sys


class Tree:
    def __init__(self, sum, range):
        self.sum = sum
        self.range = range
        self.left = None
        self.right = None


def makeTree(node, arr):
    if node.range[0] > node.range[1]:
        return 0
    if node.range[0] == node.range[1]:
        node.sum = arr[node.range[0]]
        return node.sum
    mid = int((node.range[0] + node.range[1])/2)
    if node.range[0] <= mid and mid <= node.range[1]:
        leftNode = Tree(0, [node.range[0], mid])
        rightNode = Tree(0, [mid+1, node.range[1]])
        node.left = leftNode
        node.right = rightNode
        node.sum = makeTree(leftNode, arr) + makeTree(rightNode, arr)
        return node.sum
    else:
        return 0


def getSum(node, start, end):
    if node.range[1] < start or node.range[0] > end:
        return 0
    elif start <= node.range[0] and node.range[1] <= end:
        return node.sum
    sum = getSum(node.left, start, end) + getSum(node.right, start, end)
    return sum


def update(node, index, diff):
    if index < node.range[0] or node.range[1] < index:
        return
    if node.range[0] <= index and index <= node.range[1]:
        node.sum += diff
    if node.left != None and node.right != None:
        update(node.left, index, diff)
        update(node.right, index, diff)


N, M, K = map(int, sys.stdin.readline().split())
arr = [0] * (N+1)
for i in range(1, N+1):
    num = int(input())
    arr[i] = num

root = Tree(0, [1, N])
makeTree(root, arr)

for _ in range(M + K):
    a, b, c = map(int, sys.stdin.readline().split())
    if a == 1:
        diff = c - arr[b]
        arr[b] = c
        update(root, b, diff)
    else:
        print(getSum(root, b, c))