[BAEKJOON] 2042. 구간 합 구하기

Posted on October 04, 2020

문제 :

어떤 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보다 크거나 같고, 2631보다 작거나 같은 정수이다.

출력 :

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

풀이 :

세그먼트 트리(Segment Tree)를 사용해야 풀 수 있는 문제이다. 다른 공부를 한다는 핑계로 세그먼트 트리를 공부하는 것을 미뤄왔는데 추석 연휴가 가기 전에 맛은 봐야되지 않겠나 싶은 생각이 들어 세그먼트 트리에 대해 알아보았다.

우선 본 문제의 경우 주어지는 숫자가 최대 100만 개일 뿐 더러, 구간합을 찾아야 하는 횟수도 만 개 이하로 완전 탐색으로는 풀 수 없는 문제이다. 때문에 무조건 세그먼트 트리를 활용하라고 만든 문제인 것 같았다. (물론, 다른 방법으로 푸는 분들도 있겠지..?)

우선 세그먼트 트리의 경우 초반 구간합을 계산할 트리를 세팅하고 구간합을 계산하는 sum 함수를 만들어야 한다. 또 문제에서는 배열의 숫자를 변경하는 작업도 수행해야 하기 때문에 배열의 숫자를 변경하고 구간합도 모두 수정하는 과정의 update 함수도 필요하다. 세그먼트 트리를 충실히 공부한다면 크게 어렵지 않은 문제인 듯 하나 세그먼트 트리가 처음인지라 공부하며 문제를 풀어 시간이 꽤나 걸렸다.

이번 기회로 세그먼트 트리 자료 구조를 완벽히 이해하고 더 많은 문제를 풀어봐야겠다.

코드 :

코드 보기/접기

#include <iostream>
#include <cmath>
#define ll long long

using namespace std;
ll* tree, * input;
int n;

ll makeTree(int start, int end, int index) {
    if (start == end) {
        tree[index] = input[end];
        return tree[index];
    }

    int mid = (start + end) / 2;
    tree[index] = makeTree(start, mid, index * 2) + makeTree(mid + 1, end, index * 2 + 1);
    return tree[index];
}

ll partialSum(ll start, ll end, int curl, int curr, int index) {
    if (start <= curl && curr <= end)   return tree[index];
    if (curr < start || end < curl) return 0;

    int mid = (curl + curr) / 2;
    return partialSum(start, end, curl, mid, index * 2) + partialSum(start, end, mid + 1, curr, index * 2 + 1);
}

void updateNum(ll index, ll num, int start, int end, int cur) {
    if (start == end && end == index) {
        tree[cur] = num;
        return;
    }

    int mid = (start + end) / 2;
    if (mid < index) updateNum(index, num, mid + 1, end, cur * 2 + 1);
    else updateNum(index, num, start, mid, cur * 2);
    tree[cur] = tree[cur * 2] + tree[cur * 2 + 1];
}

int main() {
    ios_base::sync_with_stdio(false);  cin.tie(NULL);  cout.tie(NULL);
    int flag, i, m, k;
    ll b, c;
    cin >> n >> m >> k;
    m += k;
    input = new ll[n + 1];
    for (i = 1; i <= n; i++)
        cin >> input[i];
    tree = new ll[int(pow(2, ceil(log(n) / log(2)) + 1))];
    makeTree(1, n, 1);
    while (m--) {
        cin >> flag >> b >> c;
        if (flag == 1) updateNum(b, c, 1, n, 1);
        else cout << partialSum(b, c, 1, n, 1) << '\n';
    }
    return 0;
}

Algorithm

algorithm baekjoon C++