백준 2042 구간 합 구하기 (파이썬) (세그먼트)

풀이

import sys; from math import ceil, log

def main(node, start, end):
    if start == end:
        tree[node] = nums[start]
        return tree[node]
    else:
        tree[node] = main(node*2, start, (start+end)//2) + main(node*2+1, (start+end)//2+1, end)
        return tree[node]

def sub_sum(node, start, end, left, right):
    if start > right or end < left:
        return 0

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

    return sub_sum(node*2, start, (start+end)//2, left, right) + sub_sum(node*2+1, (start+end)//2+1, end, left, right)


# 요소가 바뀌는 경우의 구간 합 구하기 -> 세그먼트 트리 사용 이유
# 일단 start와 end가 움직이며 값을 바꿀 것,
def update(node, start, end, idx, val):
    # 이 범위에 idx 없으면 return
    if start > idx or idx > end:
        return
    # return 안했으면 일단 값 바꿔주고
    tree[node] += val

    # 리프노드까지 계속 들어가야됨
    if start != end:
        update(node*2, start, (start+end)//2, idx, val)
        update(node*2+1, (start+end)//2+1, end, idx, val)

if __name__ == '__main__':
    N, M, K = map(int, sys.stdin.readline().strip().split())
    nums = []
    for _ in range(N):
        nums.append(int(sys.stdin.readline().strip()))

    tree = [0] * 2**(ceil(log(N, 2)) + 1)
    main(1, 0, N-1)
    for _ in range(M+K):
        a, b, c = map(int, sys.stdin.readline().strip().split())
        if a == 1:
            # 기존 값과의 차이
            diff = c-nums[b-1]
            # 기존에 있는 값을 바꿔주고
            nums[b-1] = c
            # 트리 업데이트
            update(1, 0, N-1, b-1, diff)
        elif a == 2:
            print(sub_sum(1, 0, N-1, b-1, c-1))

 

---

후기

• 이번 코드는 저번 구간 합 구하기 4에서 update함수만 추가되어서 이 부분만 설명할게요
• ↓ 구간 합 구하기 4 (세그먼트 트리의 생성 및 구간 합 구하는 코드)
• https://ryuwc.tistory.com/193
• 이번에는 위의 링크에서 이어서 update함수만 만들어주면 된다.
• update할 때에는 일단 최상단 1번 노드부터 바꾸며 시작 (요소 한개가 바뀌면 일단 전체 합도 바뀜)
• 그리고 점점 쪼개서 들어가서 idx에서 벗어나면 return 하고 idx가 범위에 포함되면 바꿔
• 위의 과정을 리프노드까지 진행
• 그림으로 그리면서 설명하면 편한데, 태블릿이 없어서 설명 못함
• 직접 그려보며 코드를 짜보시오