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

코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>
#include <cmath>

using namespace std;
typedef long long ll;


ll tc, n, m, a[1000004],x,y,z,k;

void init(vector<ll> &tree, int node, int start, int end) {
    // 리프노드
    if(start == end) {
        tree[node] = a[start];
    } else {
        int mid = (start + end) / 2;
        init(tree, node*2, start, mid); // 부모 기준 왼쪽 노드 탐색
        init(tree, node*2 + 1, mid + 1, end); // 부모 기준 오른쪽 노드 탐색
        tree[node] = tree[node*2] + tree[node*2 + 1];
    }
}

ll query(vector<ll> &tree, int node, int start, int end, int left, int right) {
    if(start > right || end < left) return 0; // 범위를 벗어난 경우
    if(start >= left && end <= right) return tree[node]; // 범위에 속한 경우
    // 일부만 속한 경우
    int mid = (start + end) / 2;
    ll lsum = query(tree, node*2, start, mid, left, right); // 왼쪽 노드 탐색
    ll rsum = query(tree, node*2 + 1, mid + 1, end, left, right); // 오른쪽 노드 탐색
    return lsum + rsum;
}

void update(vector<ll> &tree, int node, int start, int end, int idx, ll val) {
    if(idx > end || idx < start) return;
    // 리프노드
    if(start == end) {
        a[idx] = val;
        tree[node] = val;
        return;
    }
    int mid = (start + end) / 2;
    update(tree, node*2, start, mid, idx, val); // 왼쪽 노드 탐색
    update(tree, node*2 + 1, mid + 1, end, idx, val); // 오른쪽 노드 탐색
    tree[node] = tree[node*2] + tree[node*2 + 1];
}


int main() {
    cin >> n >> m >> k;
    int h = (int)ceil(log2(n));
    int hei = 2*(1 << h)-1;
    vector<ll> tree(hei + 1);
    for(int i = 0; i < n; i++) {
        cin >> a[i];
    }
    init(tree, 1, 0, n - 1);
    for(int i = 0; i < m + k; i++) {
        cin >> x >> y >> z;
        if(x == 1) {
        	// 수정
            update(tree, 1, 0, n - 1, y - 1, z);
        } else if(x == 2) {
            // 값 출력
            cout << query(tree, 1, 0, n - 1, y - 1, z - 1) << '\n';
        }
    }
    return 0;
}

세그먼트 트리를 학습할 때 기본이 되는 문제라고 한다. 배열 요소의 변경이 잦은 경우, 구간합을 구하는 것에 있어 시간 복잡도가 증가한다. 변경된 값으로 인해 구간합을 다시 구해야하기 때문이다. 세그먼트 트리를 사용하면 누적 합을 구하는 것과 수정하는 연산을 O(logN)에 수행할 수 있다고 한다. 따라서 중간에 배열 요소가 수정되는 경우 세그먼트 트리를 사용하면 빠르게 문제를 해결할 수 있을 것이다.