[C++] 백준 2042 구간 합 구하기

[!NOTE]

세그먼트 트리 자료 구조에 대한 설명을 보려면 다음 포스트 참조: 세그먼트 트리 (Segment Tree) - C/C++로 설명

세그먼트 트리는 구간 합이나 최소값, 최대값을 구하는 데 유용한 자료구조이다. 특히 데이터가 자주 업데이트되면서도 특정 구간의 값을 빠르게 조회해야 하는 문제에서 큰 효율성을 제공한다. 이번 글에서는 백준 온라인 저지의 2042번 문제, 구간 합 구하기를 통해 세그먼트 트리를 구현하고 문제를 해결하는 과정을 설명한다.

---

문제 설명

N개의 수가 주어질 때, 중간에 특정 수를 변경하면서 주어진 구간의 합을 빠르게 구하는 문제이다.

예제 입력 설명

1부터 5까지의 수가 주어지고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지의 합을 구하라고 한다면 결과는 17이 된다. 그 상태에서 5번째 수를 2로 바꾸고 3번째부터 5번째까지의 합을 구하면 결과는 12가 된다.

입력 및 출력 형식

• 입력
  • 첫째 줄에 수의 개수 N (1 ≤ N ≤ 1,000,000)과, 변경 횟수 M (1 ≤ M ≤ 10,000), 구간합 질의 횟수 K (1 ≤ K ≤ 10,000)가 주어진다.
  • 둘째 줄부터 N개의 수가 주어진다.
  • M + K 개의 줄에 세 개의 정수 a, b, c가 주어진다:
    • a = 1일 때: b번째 수를 c로 변경
    • a = 2일 때: b번째부터 c번째까지의 합을 구함.
• 출력
  • 구간합 질의에 대해 하나씩 결과를 출력한다.

---

문제 접근 및 세그먼트 트리 설계

일반적인 구간합 구하기 방식은 비효율적이다. 세그먼트 트리를 사용하면, 구간합 질의와 특정 원소의 업데이트를 $O(log N)$의 시간복잡도로 해결할 수 있다.

세그먼트 트리 구성

세그먼트 트리는 다음과 같은 세 가지 함수로 구성된다:

1. init: 초기 배열을 기반으로 세그먼트 트리를 구성하는 함수
2. update: 배열의 특정 원소가 바뀔 때, 이를 반영하여 트리를 업데이트하는 함수
3. sum: 특정 구간의 합을 구하는 함수

---

구현 코드

다음은 문제 해결을 위한 C++ 코드이다.

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#define MAX 1000000
using namespace std;
using ll = long long;

int N, M, K;
vector<ll> tree(4 * MAX), arr(MAX);

// 세그먼트 트리 초기화 함수
ll init(int node, int left, int right) {
    if (left == right) 
        return tree[node] = arr[left];

    int mid = (left + right) / 2;
    return tree[node] = init(2 * node, left, mid) + init(2 * node + 1, mid + 1, right);
}

// 세그먼트 트리 업데이트 함수
void update(int node, int left, int right, int idx, ll diff) {
    if (idx < left || idx > right) return;
    tree[node] += diff;

    if (left != right) {
        int mid = (left + right) / 2;
        update(2 * node, left, mid, idx, diff);
        update(2 * node + 1, mid + 1, right, idx, diff);
    }
}

// 구간 합 구하기 함수
ll sum(int node, int left, int right, int arrL, int arrR) {
    if (arrR < left || right < arrL) return 0;  // 범위 밖
    if (arrL <= left && right <= arrR) return tree[node]; // 범위 안

    int mid = (left + right) / 2;
    return sum(2 * node, left, mid, arrL, arrR) + sum(2 * node + 1, mid + 1, right, arrL, arrR);
}

int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    // 입력
    cin >> N >> M >> K;
    for (int i = 1; i <= N; i++)
        cin >> arr[i];

    // 세그먼트 트리 초기화
    init(1, 1, N);

    // M + K번의 연산 처리
    while (M + K > 0) {
        int cmd;
        cin >> cmd;

        if (cmd == 1) {  // update 연산
            int i; ll v;
            cin >> i >> v;
            ll diff = v - arr[i];
            update(1, 1, N, i, diff);
            arr[i] = v;
            M--;
        } 
        else if (cmd == 2) {  // 구간 합 연산
            int l, r;
            cin >> l >> r;
            cout << sum(1, 1, N, l, r) << '\n';
            K--;
        }
    }
    return 0;
}

---

코드 설명

1. init 함수 (세그먼트 트리 초기화)

init 함수는 세그먼트 트리를 구성하는 재귀 함수이다.

• 왼쪽, 오른쪽 자식 노드의 합을 부모 노드에 저장하여 트리를 완성한다.
• 배열의 리프 노드에 직접 값을 저장하고, 부모 노드에 자식 노드의 값을 더하여 저장한다.

2. update 함수 (원소 업데이트)

update 함수는 특정 위치의 원소를 변경할 때 사용하는 함수이다.

• 변경된 값이 포함된 구간만을 업데이트하기 때문에, 업데이트 시간은 O(log N)이다.
• 노드에 도달할 때마다 diff 값을 더해 전체 트리를 갱신한다.

3. sum 함수 (구간 합 계산)

sum 함수는 특정 구간의 합을 구하는 함수이다.

• 구간이 완전히 포함되면 해당 노드 값을 반환하고, 겹치지 않으면 0을 반환한다.
• O(log N)의 시간복잡도로 구간 합을 구할 수 있다.

---

예제 풀이 과정

예제 입력:

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

1. 세그먼트 트리를 초기화하여 [1, 2, 3, 4, 5]의 배열을 트리에 저장.
2. 1 3 6 명령어로 3번째 원소를 6으로 변경.
3. 2 2 5 명령어로 2번째부터 5번째 구간 합(17)을 출력.
4. 1 5 2 명령어로 5번째 원소를 2로 변경.
5. 2 3 5 명령어로 3번째부터 5번째 구간 합(12)을 출력.

출력

17
12

---

세그먼트 트리는 구간 합과 같은 범위 연산을 효율적으로 처리할 수 있는 강력한 자료구조로, 특히 값의 변경이 빈번한 경우 유용하다. 이번 예제는 세그먼트 트리의 기본 개념을 이해하고, 이를 통해 구간 합 구하기 문제를 해결하는 좋은 사례가 될 수 있다.