[Python3] 2042 구간 합 구하기 - 골드 1

• 2022.02.10 포스팅 기준 solved.ac 문제 등급 : 골드 1
• 단계별로 풀어보기 -> 세그먼트 트리
• 알고리즘 분류
  • 자료구조
  • 세그먼트 트리

 

<링크>
https://www.acmicpc.net/problem/2042

2042번: 구간 합 구하기

시간제한	메모리제한
2 초	256 MB

 

---

<문제>
어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

 

<입력>
첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데, a가 1인 경우 b(1 ≤ b ≤ N)번째 수를 c로 바꾸고 a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.
입력으로 주어지는 모든 수는 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

 

<출력>
첫째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

 

<예제 입력 1>

5 2 2
1
2
3
4
5
1 3 6
2 2 5
1 5 2
2 3 5

<예제 출력 1>

17
12

 

---

<풀이 접근>
세그먼트 트리에 대한 자세한 설명은 백준온라인저지 블로그에서 보실 수 있습니다.
https://www.acmicpc.net/blog/view/9

세그먼트 트리 (Segment Tree)

 

이 문제에서 요구하는 세그먼트 트리의 기능으로는,

1. 세그먼트 트리의 초기화
2. 트리 내의 특정 구간의 합
3. 트리의 갱신

으로, 총 3가지가 있습니다.

문제의 입력은 트리의 초기 상태가 주어진 후에 구간의 합을 물어보거나 수의 변경을 요구하므로, 주어진 초기 상태를 이용하여 세그먼트 트리를 초기화해주어야 합니다.
 
구간 합 기능은 위 설명글의 개념 그대로 구현할 것이고, 갱신 단계에서는 루트노드부터 목표 리프노드를 찾아 내려가면서 입력받은 수와 기존값의 차이만큼 값을 수정할 것입니다. 목표 구간의 밖에서 탐색을 할 때에는, 덧셈의 항등원인 0을 반환하도록 하는 점에 유의합니다.

*세그먼트 트리의 크기와 인덱싱에 대해서는 이전 포스팅 https://qodelog.tistory.com/2 을 참고 바랍니다.

import sys
import math

input = sys.stdin.readline


# N : 구간의 초기 상태
# M : 수 변경이 일어나는 횟수
# K : 구간 합을 구하는 횟수
n, m, k = map(int, input().split())

# n개 수의 상태를 저장. index 1부터 n까지 사용하기 위해 [0]을 붙여줍니다.
nums = [0] + [int(input()) for _ in range(n)] 

# arr : m + k 개의 쿼리 (a, b, c)를 입력받습니다.
arr = [[*map(int, input().split())] for _ in range(m + k)]

seg_tree = [0 for _ in range((1 << int(math.ceil(math.log2(size + 1)))+1) + 1)]


# 트리 초기화
def make_tree(tree, start, end, v):
    if start == end:
        tree[v] = nums[start]
    else:
        mid = (start + end) // 2
        tree[v] = make_tree(tree, start, mid, v * 2) + make_tree(tree, mid + 1, end, v * 2 + 1)
    return tree[v]


# 트리 구간 합 구하기
def sum_tree(tree, start, end, v, left, right):
    if end < left or right < start:
        return 0
    elif left <= start and end <= right:
        return tree[v]
    else:
        mid = (start + end) // 2
        return sum_tree(tree, start, mid, v * 2, left, right) + sum_tree(tree, mid + 1, end, v * 2 + 1, left, right)


# 트리 갱신
def update_tree(tree, start, end, v, idx, dif):
    if idx < start or idx > end:
        return
    tree[v] += dif
    if start == end:
        return
    mid = (start + end) // 2
    update_tree(tree, start, mid, v * 2, idx, dif)
    update_tree(tree, mid + 1, end, v * 2 + 1, idx, dif)
    return


make_tree(seg_tree, 1, n, 1)
res = ''

for a, b, c in arr:
    if a == 1:
        update_tree(seg_tree, 1, n, 1, b, c - nums[b])
        nums[b] = c
    elif a == 2:
        res += str(sum_tree(seg_tree, 1, n, 1, b, c)) + '\n'

print(res.rstrip())

 

위 코드 수행 결과는 다음과 같습니다.

이 문제 정답 코드들 중 상위권의 수행 시간이 500ms 안쪽으로 나오는 것을 보아, 최적화가 필요해 보입니다. 

---

<마무리>

이 문제는 구간의 합을 물어보며, 트리 값의 갱신을 요구하는 기본적인 세그먼트 트리 문제입니다.
다음에 다뤄 볼 구간의 곱을 구하는 문제, 최대값과 최솟값을 구하는 문제와 어떤 점이 다른지 비교해보면 좋을 것 같습니다.