백준 2042번 구간 합 구하기 C++ 세그먼트 트리

리뷰

세그먼트 트리로 구간 합을 구하는 문제

 

문제 풀이

1. n, m, k 값을 입력 받고 n개의 줄로 입력되는 값을 nodes 배열에 받아준다.
2. 트리를 4 * n크기로 초기화 하고 nodes 배열을 세그먼트 트리로 빌드해 준다.
3. m + k를 합친 수만큼 반복문을 실행하고 a, b, c를 입력 받는다.
4. 만약 첫번째 수가 1이라면 b - 1인덱스의 노드 트리를 v 값으로 update 해준다.
5. 첫번째 수가 2라면 query 함수를 통해 b - 1 ~ c - 1 사이의 값을 출력해 준다.

 

참고 사항

입력 값이 -2^63 ~ 2^63-1 범위 이므로 int가 아닌 long long 타입을 사용해 준다.

 

 

정답 코드

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int n, m, k;
vector<long long> tree;

void build(vector<long long>& nodes, int node, int start, int end) {
    if (start == end) {
        tree[node] = nodes[start];
    }
    else {
        int mid = (start + end) / 2;
        build(nodes, 2 * node, start, mid);
        build(nodes, 2 * node + 1, mid + 1, end);
        tree[node] = tree[2 * node] + tree[2 * node + 1];
    }
}

void update(int node, int start, int end, int indx, long long value) {
    if (start == end) {
        tree[node] = value;
    }
    else {
        int mid = (start + end) / 2;
        if (start <= indx && indx <= mid) {
            update(2 * node, start, mid, indx, value);
        }
        else {
            update(2 * node + 1, mid + 1, end, indx, value);
        }
        tree[node] = tree[2 * node] + tree[2 * node + 1];
    }
}

long long query(int node, int start, int end, int left, int right) {
    if (right < start || end < left) {
        return 0;
    }
    if (left <= start && end <= right) {
        return tree[node];
    }
    int mid = (start + end) / 2;
    long long left_min = query(2 * node, start, mid, left, right);
    long long right_min = query(2 * node + 1, mid + 1, end, left, right);
    return left_min + right_min;
}

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

    cin >> n >> m >> k;
    vector<long long> nodes;
    tree.resize(4 * n);
    
    for (int i = 0; i < n; i++) {
        long long temp;
        cin >> temp;
        nodes.push_back(temp);
    }
    build(nodes, 1, 0, n - 1);

    int go = m + k;
    while (go--) {
        long long order, b, c;
        cin >> order >> b >> c;
        if (order == 2) {
            cout << query(1, 0, n - 1, b - 1, c - 1) << "\n";
        }
        else {
            update(1, 0, n - 1, b - 1, c);
        }
    }
}