[Backjoon] 골드 1 구간 합 구하기

문제: 골드 1 구간 합 구하기

https://www.acmicpc.net/problem/2042

풀이코드

import sys
input = sys.stdin.readline

n, m, k = map(int, input().split())
lst = []
for _ in range(n):
    lst.append(int(input()))

## 리프 노드 개수 설정 -> n 이상의 2 제곱수
size = 1
while size < n: 
    size *= 2
    
tree = [0] * 2 * size

## 리프 노드 채우기
for i in range(n): 
    tree[size + i] = lst[i]
## 내부 노드 채우기
for i in range(size - 1, 0, -1): 
    tree[i] = tree[2*i] + tree[2*i+1]

## 값 교체
def update(idx, val): 
    pos = size + idx
    tree[pos] = val
    pos //= 2
    while pos > 0: 
        tree[pos] = tree[2*pos] + tree[2*pos+1]
        pos //= 2
## 구간 합
def rangesum(start, end): 
    start_pos = size + start
    end_pos = size + end
    s = 0
    while start_pos <= end_pos: 
        if start_pos % 2 == 1: 
            s += tree[start_pos]
            start_pos += 1
        if end_pos % 2 == 0: 
            s += tree[end_pos]
            end_pos -= 1
        start_pos //= 2
        end_pos //= 2
    return s

for _ in range(m + k): 
    a, b, c = map(int, input().split())
    if a == 1: ## 교체
        update(b-1, c)
    else: ## 구간 합 구하기
        print(rangesum(b-1, c-1))

---

느낀 점

일단 이해하기가 어려워서 그려보면서 이해를 했다.

구간 합을 구하는 건 segment tree를 활용해야 했다. 여기서 처음엔 Top-down 방식으로 사용하다가 재귀를 써야 해서 조금 복잡하기도 하고 이해가 조금 부족한 것 같아서 Bottom-up 방식으로 다시 풀어보려고 한다. 

 

가장 먼저 트리의 크기를 정해야 한다.

segment tree는 완전 이진 트리 구조이며 루트 노드는 배열 전체 구간의 합이 저장되어있고, 루트의 왼쪽 자식 노드에는 앞부터 절반까지의 합, 루트의 오른쪽 자식 노드에는 뒤쪽 절반까지의 합이 저장되어 있다. 이게 계속해서 반을 쪼개가면서 리프 노드까지 내려간다. 

리프 노드의 개수는 배열 뭔소의 개수와 같다. 즉, segment tree를 배열로 표현하려면, 리프 수가 2의 제곱수여야 한다는 것이다. 

배열의 원소 개수가 2의 제곱수라면 딱 맞게 완전 이진 트리가 만들어지지만 배열의 언소 개수가 2의 제곱수가 아니라면 트리의 크기를 정하기가 어려워진다. 

때문에 배열 원소의 개수가 딱 2의 제곱이 아니라면 배열 원소의 개수보다 크거나 같은 최소 2의 제곱수로 리프 노드의 개수로 맞춰준다. 

 

문제에서 나온 예시를 들어보자. 

## 배열: [a, b, c, d, e]
## n=5
## -> n보다 크거나 같은 최소 2의 제곱수 = 8
## -> 리프는 8개 필요

## n=5, arr = [1,2,3,4,5], size=8

                          [1]=15               ← 루트 (전체 합)
            /                             \
         [2]=10                           [3]=5
    /              \                 /              \
 [4]=3            [5]=7            [6]=5           [7]=0
 /  \              /  \             /  \            /  \
[8]=1 [9]=2    [10]=3 [11]=4   [12]=5   [13]=0  [14]=0  [15]=0

## 리프: [8..12] ← arr 값 1,2,3,4,5
## 나머지 리프 [13..15]는 0

 

만약 노드의 값을 바꾼다면 어떻게 될까? 
만약 arr[2] = 3을 10으로 바꾼다고 할 때...

pos = size + idx = 8 + 2 = 10

tree[10]: 3 → 10 

그러면 자연스럽게 pos가 10 / 2 = 5인 노드의 값도 변해야 하며 5 // 2 = 2인 노드의 값, 2 / 2 = 1인 노드의 값까지 변하면 된다. 

 

만약 구간합을 구한다면 어떻게 될까? 

만약 arr에서 start = 0 부터 end = 2까지 1,2,3 의 구간합을 구한다고 해보자. 

일단 리프 노드 인덱스로 바꾸면 

start_pos = size + start = 8

end_pos = size + end = 10

만약 노드의 start_pos 가 짝수이고, end_pos가 홀수라면 부모 노드로 올라가도 된다. 그러면 start_pos //= 2, end_pos //= 2 해주고... 이를 반복한다. 이후 start_pos == end_pos이면 해당 노드의 값이 정답이 된다. 

하지만 만약 노드의 start_pos 가 홀수고 end_pos가 짝수라면 부모 노드로 올라가면 안 되기 때문에 tree[start_pos]를 더해주고 start_pos += 1 해주거나 tree[end_pos]를 더해주고 end_pos -= 1을 해준다. 

 

 

홀짝 조합별 행동을 요약하면 다음과 같다. 

 

• (짝수, 홀수) → 그냥 부모로 올라감 (아무것도 안 더함)
• (홀수, 짝수) → 둘 다 더하고 부모로 올라감
• (홀수, 홀수) → start_pos 노드만 더하고 부모로 올라감
• (짝수, 짝수) → end_pos 노드만 더하고 부모로 올라감