[백준 2042번] 구간 합 구하기 (C++, Python)

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

2042번: 구간 합 구하기

이번 포스트에서는 백준 알고리즘 문제인 구간 합 구하기를 포스트 하겠습니다.

 

해당 문제는 solved.ac 기준 골드 1에 위치한 문제입니다.

 

먼저 문제입니다.

수의 개수 N, 수의 변경이 일어나는 횟수 M, 구간의 합을 구하는 횟수 K를 입력받습니다.

그리고, N개의 수를 입력받고, M+K 개의 구간합 혹은 수의 변경 명령을 입력받습니다.

이때, 구간합을 구하는 경우 구간의 합을 구하여 출력하는 문제입니다.

 

해당 문제는 시간 제한이 2초에 N이 최대 100만으로 주어집니다. 필자는 그래서 선형으로 해결을 하려 했으나, 시간초과가 발생합니다.

 

O(N)의 시간복잡도로 구간합을 구하게 되면 M+K개의 명령이 전부 구간합 이고, 1~N까지의 합을 요구하게 되면 시간이 오래 걸리게 되어 시간 초과가 발생합니다. ( 수의 변경은 O(1) )

 

그러므로 해당 문제는 세그먼트 트리로 해결합니다.

 

먼저 세그먼트 트리는 트리의 각 원소가 트리의 구간 정보를 저장하고 있어서 위 문제 기준으로 구간 합, 값 변경을 모두 O(log N)의 시간복잡도로 수행하게 됩니다.

 

세그먼트 트리는 Init() 함수와 update(), subSum() 함수로 구성됩니다.

 

Init() 함수는 세그먼트 트리의 구간을 정의하고, 초기값을 설정합니다.

리프 노드에 도달할 때 까지 재귀적으로 자식 노드로 접근하고, 리프 노드에 도달하게 되면 해당 구간 값을 리프 노드에 할당합니다. 그리고 리턴된 값을 더해서 부모 노드에 할당합니다. (이진 트리이기 때문에 자식 노드 2개의 구간 합을 부모 노드에 저장하는 원리)

 

update() 함수는 값을 변경하는 함수로, 해당 구간이 포함된 모든 구간의 원소들을 갱신합니다.

자식 노드들에 재귀적으로 접근하며 변경된 차이값을 해당 구간에 더해줍니다.

 

subSub() 함수는 구해놓은 구간 값들의 합으로 구간합을 리턴합니다.

똑같이 재귀적으로 자식노드에 접근하고, 현재 접근한 노드가 구간합을 구하고자 하는 범위에 포함되면 그대로 그 노드를 리턴합니다. 그리고 리턴받은 값의 합을 리턴합니다.

 

다음은 C++ 코드입니다.

 

#include <iostream>
 
using namespace std;
 
long long int arr[1000001];
long long int segTree[4000004];
 
long long int Init(int start, int last, int nodeNum) {
    if(start == last)
        return segTree[nodeNum] = arr[start];
    
    int mid = (start + last) / 2;
    return segTree[nodeNum] = Init(start, mid, nodeNum * 2) + Init(mid + 1, last, nodeNum * 2 + 1);
}
 
long long int subSum(int start, int last, int nodeNum, int left, int right) {
    if(left > last || right < start)
        return 0;
 
    if(left <= start && right >= last) {
        return segTree[nodeNum];
    int mid = (start + last) / 2;
    return subSum(start, mid, nodeNum * 2, left, right) + subSum(mid + 1, last, nodeNum * 2 + 1, left, right);
}
 
void update(int start, int last, int nodeNum, int targetIdx, long long int val) {
    if(targetIdx < start || targetIdx > last)
        return;
 
    segTree[nodeNum] += val;
    if(start == end)
        return;
    int mid = (start + last) / 2;
    update(start, mid, nodeNum * 2, targetIdx, val);
    update(mid + 1, last, nodeNum * 2 + 1, targetIdx, val);
}
 
int main() {
    int n, m, k;
    cin >> n >> m >> k;
 
    for(int i = 0; i < n; i++) {
        cin >> arr[i];
    }
 
    Init(0, n - 1, 1);
 
    long long int a,b,c;
    for(int i = 0; i < m + k; i++) {
        cin >> a >> b >> c;
        if(a == 1) {
            long long int tmp = c - arr[b - 1];
            arr[b - 1] = c;
            update(0, n - 1, 1, b - 1, tmp);
        } else {
            cout << subSum(0, n - 1, 1, b - 1, c - 1) << "\n";
        }
    }
 
    return 0;
}

다음은 Python 코드입니다.

 

arr = [0] * 1000001
segTree = [0] * 4000004
 
def Init(start, last, nodeNum) :
    if start == last :
        segTree[nodeNum] = arr[start]
        return segTree[nodeNum]
 
    mid = (start + last) // 2
    segTree[nodeNum] = Init(start, mid, nodeNum * 2) + Init(mid + 1, last, nodeNum * 2 + 1)
    return segTree[nodeNum]
 
def subSub(start, last, nodeNum, left, right) :
    if left > last or right < start :
        return 0
 
    if left <= start and right >= last :
        return segTree[nodeNum]
    mid = (start + last) // 2
    return subSum(start, mid, nodeNum * 2, left, right) + subSum(mid + 1, last, nodeNum * 2 + 1, left, right)
 
def update(start, last, nodeNum, targetIdx, val) :
    if targetIdx < start or targetIdx > last :
        return
 
    segTree[nodeNum] += val
    if start == end :
        return
    mid = (start + last) // 2
    update(start, mid, nodeNum * 2, targetIdx, val)
    update(mid + 1, last, nodeNum * 2 + 1, targetIdx, val)
}
 
n,m,k = map(int, input().split())
 
for i in range(n) :
    arr[i] = int(input())
 
Init(0, n - 1, 1)
 
for i in range(m + k) :
    a,b,c = map(int, input().split())
    if a == 1 :
        tmp = c - arr[b - 1]
        arr[b - 1] = c
        update(0, n - 1, 1, b - 1, tmp)
    else :
        print(subSum(0, n - 1, 1, b - 1, c -1))

해당 문제는 세그먼트 트리를 이해하고 Init(), update(), subSum() 함수를 작성할 수 있어야 해결할 수 있는 문제로 세그먼트 트리를 입문하기 괜찮은 문제라고 생각합니다. 세그먼트 트리 문제들은 기본 골드1 ~ 플레티넘 이므로 난이도가 높지만 도전해보기 재밌는 문제입니다.

 

중간에 시간초과, 런타임 에러는 선형으로도 작성하여 실험해본 결과입니다.

 

파이썬의 경우 import sys, input = sys.stdin.readline을 사용하여야 시간초과가 발생하지 않습니다.

 

다음 포스트에서도 다른 백준 알고리즘 문제로 포스트하도록 하겠습니다. 감사합니다.