[ BOJ - 2042번 ] 구간 합 구하기 (C++)

https://www.acmicpc.net/problem/2042

 

 배열의 크기가 최대 1000000이므로, 배열의 구간 합에 대한 Range query와 배열의 한 원소에 대한 Point update가 O(n log n) 시간 안에 처리되어야 한다.

 따라서 Segment Tree 또는 Binary Indexed Tree를 사용하여 풀 수 있는 문제로 볼 수 있다. 원소의 절대값의 크기가 최대 | 2⁶³-1 | 이므로 long long 자료형을 사용해 원소와 구간 합 결과를 저장해야 한다는 점을 주의해야 한다.

 

[ 풀이 방법 ]

세그먼트 트리(Segment Tree)

 Segment Tree를 구현하고, 문제에서 주어진 조건에 맞게 구간 합과 원소 업데이트를 수행하여 문제를 해결하였다.

 

[ 정답 코드 ]

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

using namespace std;

typedef long long ll;

struct SegmentTree {
    int n;
    int e;
    vector<ll> tree;
    vector<ll> arr;
    
    void init(vector<ll>& origin) {
        arr.assign(origin.begin(), origin.end());
        n = int(arr.size());
        e = 0;
        tree.resize(4*n);
        
        build(1, 0, n-1);
    }
    
    ll merge(ll leftValue, ll rightValue) {
        return leftValue + rightValue;
    }
    
    ll build(ll node, ll nodeLeft, ll nodeRight) {
        if(nodeLeft == nodeRight) {
            return tree[node] = arr[nodeLeft];
        }
        
        ll mid = nodeLeft - (nodeLeft - nodeRight)/2;
        ll leftValue = build(node*2, nodeLeft, mid);
        ll rightValue = build(node*2+1, mid+1, nodeRight);
        
        return tree[node] = merge(leftValue, rightValue);
    }
    
    ll query(ll left, ll right) {
        return queryRec(1, left, right, 0, n-1);
    }
    
    ll queryRec(ll node, ll left, ll right, ll nodeLeft, ll nodeRight) {
        
        if(left <= nodeLeft && nodeRight <= right) {
            return tree[node];
        }
        
        if(nodeRight < left || right < nodeLeft) {
            return e;
        }
        
        ll mid = nodeLeft - (nodeLeft - nodeRight)/2;
        ll leftValue = queryRec(node*2, left, right, nodeLeft, mid);
        ll rightValue = queryRec(node*2+1, left, right, mid+1, nodeRight);
        
        return merge(leftValue, rightValue);
    }
    
    void update(ll idx, ll newVal) {
        updateRec(1, idx, newVal, 0, n-1);
    }
    
    ll updateRec(ll node, ll idx, ll newVal, ll nodeLeft, ll nodeRight) {
        
        if(nodeRight < idx || idx < nodeLeft) {
            return tree[node];
        }
        
        if(nodeLeft == nodeRight) {
            return tree[node] = newVal;
        }
        
        ll mid = nodeLeft - (nodeLeft - nodeRight)/2;
        ll leftValue = updateRec(node*2, idx, newVal, nodeLeft, mid);
        ll rightValue = updateRec(node*2+1, idx, newVal, mid+1, nodeRight);
        
        return tree[node] = merge(leftValue, rightValue);
    }
};

int main() {
    
    iostream::sync_with_stdio(0);
    cin.tie(0);
    
    ll n, m, k;
    cin >> n >> m >> k;
    
    vector<ll> arr(n);
    for(int i = 0; i < n; i++) {
        cin >> arr[i];
    }
    
    SegmentTree st;
    st.init(arr);
    
    for(ll i = 0; i < m+k; i++) {
        ll a, b, c;
        cin >> a >> b >> c;
        
        if(a == 1) {
            st.update(b-1, c);
        } else {
            cout << st.query(b-1, c-1) << '\n';
        }
    }
    
    return 0;
}

 

 

Binary Indexed Tree(펜윅 트리)를 구현해서 풀 수도 있다.

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

using namespace std;

struct FT {
    vector<long long> tree;
    
    void init(long long n) {
        tree.resize(n*2);
    }
    
    void add(long long pos, long long val) {
        pos++;
        
        while(pos < tree.size()){
            tree[pos] += val;
            pos += pos & -pos;
        }
    }
    
    long long sum(long long pos) {
        long long res = 0;
        
        pos++;
        while(pos > 0) {
            res += tree[pos];
            pos &= pos-1;
        }
        return res;
    }
    
    long long sumRange(long long left, long long right) {
        long long res = sum(right);
        if(left > 0) res -= sum(left-1);
        return res;
    }
};

int main() {
    
    iostream::sync_with_stdio(false);
    cin.tie(NULL);
    
    long long n, m, k;
    FT ft;
    
    cin >> n >> m >> k;
    vector<long long> values(n);
    ft.init(n);
    
    for(long long i = 0; i < n; i++) {
        long long t;
        cin >> t;
        values[i] = t;
        ft.add(i, t);
    }
    
    for(long long i = 0; i < m+k; i++) {
        long long a, b, c;
        cin >> a;
        if (a == 1) {
            cin >> b >> c;
            ft.add(b-1, c - values[b-1]);
            values[b-1] = c;
        } else {
            cin >> b >> c;
            cout << ft.sumRange(b-1, c-1) << '\n';
        }
    }
        
    return 0;
}