구간 합 구하기

백준 2042번: 구간 합 구하기

아이디어

세그먼트 트리 연습 문제다.

루트에는 전부 합한 값이, 그 자식은 구간을 반씩 잘라서 그 구간에 해당하는 값을 가지고 있다.
왼쪽 자식은 부모 인덱스 * 2, 오른쪽 자식은 부모 인덱스 * 2 + 1이다. (루트는 1부터)

어떤 구간의 합이라도 O(logN)으로 구할 수 있다. 루트부터 필요한 구간이 나올 때까지 왼쪽, 오른쪽 재귀호출한다.

또 이렇게 합을 저장해놓으면 중간에 값의 변경이 일어나도 O(logN)으로 업데이트를 할 수 있다.
예를 들어 arr[2]의 값을 바꾸면 0~4, 0~2, 2~2에 적힌 값만 바꾸면 된다.
만약 평범하게 구간합 배열을 선언했다면 업데이트 하는데 O(N)이 걸린다.

코드

#include <bits/stdc++.h>

using namespace std;

int N, M, K;
vector<long long> arr; // 0부터 시작
vector<long long> tree; // 1부터 시작

long long init(int node, int start, int end) {
    if (start == end) {
        return tree[node] = arr[start];
    }
    else {
        return tree[node] = init(node*2, start, (start+end)/2) + init(node*2+1, (start+end)/2+1, end);
    }
}

long long sum(int node, int start, int end, int left, int right) {
    if (left > end || right < start) { // s~e가 구하려는 범위에서 완전 벗어난 경우
        return 0;
    }
    if (left <= start && end <= right) { // s~e가 구하려는 범위에 완전 포함된 경우
        return tree[node];
    }
    // 나머지
    return sum(node*2, start, (start+end)/2, left, right) + sum(node*2+1, (start+end)/2+1, end, left, right);
}

void update(int node, int start, int end, int idx, long long diff) {
    if (idx < start || idx > end) return;
    tree[node] += diff;
    if (start != end) {
        update(node*2, start, (start+end)/2, idx, diff);
        update(node*2+1, (start+end)/2+1, end, idx, diff);
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> N >> M >> K;
    for (int i = 0; i < N; i++) {
        int x;
        cin >> x;
        arr.push_back(x);
    }

    int h = (int)ceil(log2(N));
    int tree_size = (1<<h+1);
    tree.resize(tree_size);
    init(1, 0, N-1);

    for (int i = 0; i < M+K; i++) {
        long long a, b, c;
        cin >> a >> b >> c;
        if (a == 1) {
            update(1, 0, N-1, b-1, c-arr[b-1]);
            arr[b-1] = c;
        }
        else if (a == 2) {
            cout << sum(1, 0, N-1, b-1, c-1) << '\n';
        }
    }
    return 0;
}

여담

어떻게 이런 신기한걸 만들어냈을까? 역시 세상에는 똑똑한 사람들이 많다.