본문 바로가기

ALGORITHM/OTHER

[c++] 2042 구간 합 구하기 (세그먼트 트리)

세그먼트 트리 이해

https://www.acmicpc.net/blog/view/9

 

세그먼트 트리 (Segment Tree)

문제 배열 A가 있고, 여기서 다음과 같은 두 연산을 수행해야하는 문제를 생각해봅시다. 구간 l, r (l ≤ r)이 주어졌을 때, A[l] + A[l+1] + ... + A[r-1] + A[r]을 구해서 출력하기 i번째 수를 v로 바꾸기. A[i

www.acmicpc.net

 

전체 코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <cmath>
using namespace std;

typedef long long ll;

ll N, M, K, input, a, b, c;

ll init(vector<ll> &arr, vector<ll> &tree, ll node, ll start, ll end) {
    if (start == end) {
        return tree[node] = arr[start];
    }
    else {
        return tree[node] = init(arr, tree, node*2, start, (start+end)/2) + init(arr, tree, node*2+1, (start+end)/2+1, end);
    }
}

ll sum(vector<ll> &tree, ll node, ll start, ll end, ll left, ll right) {
    if (end < left || start > right) {
        return 0;
    }
    if (left <= start && end <= right) {
        return tree[node];
    }
    return sum(tree, node*2, start, (start+end)/2, left, right) + sum(tree, node*2+1, (start+end)/2+1, end, left, right);
}

void update(vector<ll> &tree, ll node, ll start, ll end, ll index, ll diff) {
    if (index < start || end < index) {
        return;
    }
    tree[node] += diff;
    if (start != end) {
        update(tree, node*2, start, (start+end)/2, index, diff);
        update(tree, node*2+1, (start+end)/2+1, end, index, diff);
    }
}

int main(int argc, const char * argv[]) {
    ios_base::sync_with_stdio(0);
    cin.tie(0);  cout.tie(0);
    
    cin >> N >> M >> K;
    
    vector<ll> arr(N);
    int h = (int)ceil(log2(N));
    int treeSize = (1 << (h+1));
    vector<ll> tree(treeSize);
    
    for (int i=0; i<N; i++) {
        cin >> arr[i];
    }
    
    init(arr, tree, 1, 0, N-1);
    
    for (int i=0; i<M+K; i++) {
        cin >> a >> b >> c;
        if (a==1) {
            ll diff = c - arr[b-1];
            arr[b-1] = c;
            update(tree, 1, 0, N-1, b-1, diff);
        }
        else if (a==2) {
            cout << sum(tree, 1, 0, N-1, b-1, c-1) << "\n";
        }
    }
    
    return 0;
}