[C++ 백준 풀이] 1275번 : 커피숍2 (Segment Tree, 세그먼트 트리)

이 포스트는 프로그래밍 문제 사이트 백준 Online Judge의 1275번 : '커피숍2' 문제에 대한 풀이 코드와 해설을 다루고 있습니다.

 

문제의 난이도는 Solved.ac 기준 Gold I에 해당하며, 이 문제를 풀이하기 위해 세그먼트 트리 응용 개념을 다룰 것입니다.

 

 

1275번 : 커피숍2

 

N개의 수를 입력 받고, Q개의 턴이 주어질 때 각 턴마다 x번째 수부터 y번째 수까지의 합을 구하고, a번째 수를 b로 바꾸는 과정을 반복할 때 합을 계속해서 출력하는 문제입니다.

큰 수의 범위에서 부분합과 수의 변경을 반복하는 과정을 거치기 때문에 데이터 구조로 세그먼트 트리가 적절합니다.

세그먼트 트리를 이용하여 위의 쿼리들을 처리하는 구현을 해보도록 하겠습니다.

 

 

↓ 위의 이미지와 동일한 풀이 코드는 아래 접은 글에 정리되어 있습니다.

#include <cstdio>
#include <vector>
using namespace std;

vector<int> Arr;
vector<long long> Tree;

long long initTree(int Begin, int End, int Node) {
    if(Begin == End) return Tree[Node] = Arr[Begin];
    int Mid = (Begin+End)/2;
    long long leftKey = initTree(Begin, Mid, Node*2);
    long long rightKey = initTree(Mid+1, End, Node*2+1);
    return Tree[Node] = leftKey + rightKey;
}

long long calcSum(int Begin, int End, int Node, int Left, int Right) {
    if(Left > End || Right < Begin) return 0;
    if(Left <= Begin && Right >= End) return Tree[Node];
    int Mid = (Begin+End)/2;
    long long LeftKey = calcSum(Begin, (Begin+End)/2, Node*2, Left, Right);
    long long RightKey = calcSum((Begin+End)/2+1, End, Node*2+1, Left, Right);
    return LeftKey + RightKey;
}

void updateTree(int Begin, int End, int Node, int Index, long long Diff) {
    if(Index < Begin || Index > End) return;
    Tree[Node] += Diff;
    int Mid = (Begin+End)/2;
    if(Begin < End) {
        updateTree(Begin, Mid, Node*2, Index, Diff);
        updateTree(Mid+1, End, Node*2+1, Index, Diff);
    }
}

int main() {
    int N, M, x, y, a, b;
    scanf("%d %d", &N, &M);
    Arr.resize(N);
    for(int i=0; i<N; i++) scanf("%d", &Arr[i]);
    Tree.resize(N*4);
    initTree(0, N-1, 1);
    for(int i=0; i<M; i++) {
        scanf("%d %d %d %d", &x, &y, &a, &b);
        if(x > y) { int temp = x; x = y; y = temp; }
        printf("%lld\n", calcSum(0, N-1, 1, x-1, y-1));
        long long Diff = (long long)b - (long long)Arr[a-1];
        Arr[a-1] = b;
        updateTree(0, N-1, 1, a-1, Diff);
    }
}

 

저는 이 문제를 풀이할 때 Arr과 Tree를 모두 선언하여 Arr에는 초기 입력을 받고, 이를 이용하여 Tree를 초기화한 뒤 쿼리를 처리하였습니다.

initTree 함수가 Arr의 데이터들을 모두 집어넣으며 세그먼트 트리를 초기화하는 함수입니다.

초기화 할 때 leftTree와 rightTree의 sum을 재귀적으로 합쳐서 저장하도록 합니다.

이후 x와 y를 입력받는데, 문제에서 주의할 점은 x가 y보다 큰 경우도 있다는 것입니다.

이 경우에는 x가 작은 index가 되도록 순서를 바꾸어주고 calcSum 함수를 수행합니다.

calcSum 함수는 구간 합을 구하는 함수로, 이미 구해둔 sum 값을 구간만 잡아서 세그먼트 트리에서 합쳐오면 됩니다.

이후 a와 b를 입력받아 a번째 있는 수를 b로 바꾸어야 하는데, 이 때 Diff 값을 계산해서 노드를 따라 내려가면서 Diff 값을 더해주는 수행을 해줍니다.

이 때 a번째 수는 Arr 배열에서 a-1칸에 저장되어 있음을 인지해야 합니다.

 

또 주의할 점은 오버플로우인데, Diff를 계산하는 과정에서 간단한 예시로 20억 - (-20억)을 연산하면 int 형에서는 overflow가 발생하게 됩니다.

이러한 문제를 애초에 발생시키고 싶지 않으면 모든 데이터형을 long long으로 사용하면 되긴 하지만, 알고리즘을 공부하는 입장이기 때문에 최대한 적절한 데이터타입을 사용하여 문제를 해결하였습니다.

 

 

위의 코드를 제출해보면 160ms의 시간을 소모하고 모든 테스트케이스를 통과함을 알 수 있습니다.

x와 y의 크기가 반대로 입력될 수 있다는 점과, Diff 연산에서의 overflow만 주의하면 쉽게 해결이 가능한 문제였습니다.