백준 2042 구간 합 구하기 (C++)

2042번: 구간 합 구하기

세그먼트 트리를 이용한 문제이다. 구간 합을 구하는 방식을 세그먼트 트리를 이용하여 O(logN)으로 합을 구할 수 있다. 먼저 입력받은 값을 이용해 세그먼트 트리를 만들어준다. 그리고 1일 경우 update를 해주고 2일 경우 sum을 이용해 구간 합을 구해 출력해주었다.
세그먼트 트리를 이용하는 문제를 처음 풀어봤다. 세그먼트 트리에 대해서는 개념 정도만 공부를 하면서 알고 있었는데 처음으로 직접 구현을 해보면서 원리를 제대로 알 수 있었다. 세그먼트 트리에 대해 잘 알아두어야 겠다.

#include <iostream>
#include <vector>

using namespace std;

int N, K, M;
long long A[1000001];
long long a, b, c;
vector<long long> tree;

long long createTree(int start, int end, int idx) {
    if (start == end) return tree[idx] = A[start];

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

long long sum(int start, int end, int left, int right, int idx) {
    if (right < start || left > end) return 0;
    if (left <= start && right >= end) return tree[idx];

    int mid = (start + end) / 2;
    return sum(start, mid, left, right, idx * 2) + sum(mid + 1, end, left, right, idx * 2 + 1);
}

void update(int start, int end, int idx, int change, long long dif) {
    if (change < start || change > end) return;

    tree[idx] += dif;

    if (start == end) return;

    int mid = (start + end) / 2;
    update(start, mid, idx * 2, change, dif);
    update(mid + 1, end, idx * 2 + 1, change, dif);
}

void solution() {
    if (a == 1) {
        update(1, N, 1, b, c - A[b]);
        A[b] = c;
    }
    else cout << sum(1, N, b, c, 1) << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL); cout.tie(NULL);

    cin >> N >> K >> M;

    for (int i = 1; i <= N; i++) {
        cin >> A[i];
    }

    tree.resize(N * 4);
    createTree(1, N, 1);

    for (int i = 0; i < K + M; i++) {
        cin >> a >> b >> c;
        solution();
    }

    return 0;
}