[백준] [C++] 1275번 커피숍2 (segment tree)

1. 문제

1275번: 커피숍2 (acmicpc.net)

1275번: 커피숍2

2. 풀이

이 문제의 핵심은 수를 업데이트 시켜줘야한다는것이다.

배열을 사용하여 i번째까지의 합을 저장하여 psum[i] - psum[j] 로 구간합을 구하는 방식은

수를 업데이트하려면 최악의 경우 O(N) 의 시간 복잡도를 가진다.

 

하지만 이런 구간들을 이진 검색 트리 형태로 표현하게되면

그 수가 속하는 구간들을 업데이트하는데  O(lgn) 의 시간복잡도로 업데이트할 수 있다.

 

구간트리는 다음과 같이 배열 한개로 표현할 수 있다.

루트 : 1 부모노드: i 왼쪽자식: i*2 오른쪽자식: i*2 + 1

#include <iostream>
#include <algorithm>
#include <vector>
 
using namespace std;
int N, Q;
vector<long long> number;
vector<long long> segtree;
 
 
long long initial(int lo, int hi, int idx) {
    if(lo + 1 == hi) {
        segtree[idx] = number[lo];
        return segtree[idx];
    }
    int mid = (lo + hi)/2;
    segtree[idx] = initial(lo, mid, idx*2) + initial(mid, hi, idx*2 + 1);
    return segtree[idx];
}
 
void change(int ni, int x, int lo, int hi, int idx) {
    int mid = (lo + hi)/ 2;
    segtree[idx] += (x -number[ni]);
    if(lo + 1 == hi) {
        return;
    }
    if(ni < mid) {
        return change(ni , x, lo, mid, idx*2);
    }
    else {
        return change(ni , x, mid, hi, idx*2 + 1);
    }
}
 
long long find(int start, int end, int idx, int lo, int hi) {
    if(start == lo && end == hi ) {
        return segtree[idx];
    }
    int mid = (lo + hi)/2;
    if(end <= mid) {
        return find(start, end, idx*2, lo, mid);
    }
    else if(mid <= start) {
        return find(start, end, idx*2 + 1, mid, hi);
    }
    else {
        return find(start, mid, idx*2, lo, mid) + find(mid, end, idx*2+1, mid, hi);
    }
 
}
 
int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(NULL);
    cout.tie(NULL);
    cin>>N>>Q;
    number = vector<long long>(N);
    segtree = vector<long long>(N*4);
    for(int i = 0; i < N; i++) {
        cin>>number[i];
    }
    initial(0, N, 1);
    for(int i = 0; i < Q; i++) {
        int lo, hi, ni;
        long long x;
        cin>>lo>>hi>>ni>>x;
        if(lo > hi) {
            swap(lo, hi);
        }
        cout<<find(lo - 1, hi, 1, 0, N)<<"\n"; 
        change(ni - 1, x, 0, N, 1);
        number[ni-1] = x;
    }
}

문제에서 lo > hi인 입력도 주어진다.