[백준 2042번] 구간 합 구하기

문제 설명

어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

입력

첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데, a가 1인 경우 b번째 수를 c로 바꾸고 a가 2인 경우에는 b번째 수부터 c번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

출력

첫째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

알고리즘

세그먼트 트리(Segment tree)

접근 방식 및 풀이

문제에 따르면 우리가 구현해야할 부분은 크게 두개이다.

1. 특정 위치의 숫자를 다른 숫자로 변경하기
2. 특정 구간의 합을 구하기

이를 평범한 방법으로 구현하게 되면, 1번 과정은 O(N)이 걸리게 되고, 2번 과정은 O(1)이 걸리게 된다. 하지만 이렇게 풀게 되면 구하는 횟수가 많아지거나 N 값이 커지게 되면 너무 많은 시간이 걸리게 된다.

그래서 사용한 것이 세그먼트 트리(Segment tree)이다. 세그먼트 트리에서 제일 끝 노드들(자식이 없는 노드)에 숫자들을 넣고 나머지 노드들에는 왼쪽 자식과 오른쪽 자식의 합을 넣는다. 그렇게 하면 1번 과정과 2번 과정 모두 O(lgN)이 걸리게 걸리게 된다. 이렇게 되면 이전보다 1번 과정에서 시간을 많이 줄일 수 있다.

이미 숫자들을 전부 리프 노드에 넣었고 나머지 노드들에는 왼쪽 자식과 오른쪽 자식의 합을 넣는 과정을 완료한 상태라고 생각하고 1번, 2번을 구현해보자.

특정 위치의 숫자를 다른 숫자로 변경하기

먼저 바꿀 리프 노드를 방문해서 원래 있던 숫자를 빼고 새로 들어갈 숫자를 넣어준다. 거기서 끝이 아니다. 자식이 변경되었기 때문에 해당 노드의 부모 노드도 변경해줘야한다. 부모 노드에서도 리프 노드에 원래 있던 숫자 만큼을 빼고 새로 들어간 숫자 만큼을 더해준다. 이렇게 하면 해당 노드가 양쪽 자식 노드의 합을 유지한다. 이런 식으로 부모가 없는 루트 노드까지 올라가면서 반복하면 된다.

특정 구간의 합을 구하기

시작과 끝 두 지점에서 부모 노드로 올라가면서 계산한다. 그리고 두 지점이 만나면 계산을 종료한다.

시작 지점에서 시작하는 계산에 경우를 먼저 생각해보자. 현재 노드가 자신의 부모의 좌측 자식일 수도 있고 우측 자식일 수도 있다. 만약 현재 노드가 좌측 자식이면 우측 자식은 우리가 구하려는 구간에 포함된다. 끝 지점이 무조건 시작 지점보다 오른쪽에 있기 때문이다. 이렇게 좌측 자식과 우측 자식의 합이 필요하면 일일이 더할 필요가 없다. 어차피 부모에 해당 값이 저장되어 있기 때문에 부모로 한칸 올라가면 된다. 하지만 만약 현재 노드가 우측 자식이면 좌측 자식은 우리가 구하려는 구간 밖이다. 따라서 부모 노드는 좌측 자식 값까지 포함하고 있기 때문에 부모 노드로 올라가지 않는다. 현재 노드 값만 더한 후, 한 칸 우측으로 넘어가서 해당 노드의 부모 노드로 올라간다.

끝 지점에서 시작하는 계산은 시작 지점에서의 계산에서 좌우측만 바꾸면 된다.

이런 식으로 올라가다가 두 지점이 만날 때 제시된 구간의 숫자들을 전부 더한 상태이기 때문에 출력하면 된다.

코드

문제 풀이에 사용한 언어: C

#include <stdio.h>
#include <stdlib.h>

int
main(){
    long long * tree;
    long long n, m, k; // 숫자 개수, 변경 횟수, 구간 합을 구하는 횟수
    long long oper; // 1이면 숫자 교체, 2이면 구간 합 구하기
    long long from, to; // 변경할 숫자의 위치, 해당 위치에 들어갈 새로운 숫자
    long long result;

    scanf("%lld %lld %lld", &n, &m, &k);
    tree = malloc(sizeof(long long) * (4 * n) ); // 메모리 제한을 생각해서 동적 할당을 사용
    
    long long base = 2*n; // 제일 아래 노드들과 상위 노드들을 나눠주는 기준점, 여유롭게 2n으로 잡음
    
    for (long long i = 1; i <= n; i++) {
        scanf("%lld", &tree[i + base]); // 제일 아래 노드에 순서대로 숫자를 저장
    }
    
    for (long long i = base; i > 0; i--) { // 밑에서 2번째 줄부터 시작해서 제일 위까지
        tree[i] = tree[i * 2] + tree[i * 2 + 1]; // 자식 노드 두개를 합한 값을 현재 노드에 저장
    }
    
    for(long long i = 0; i < m + k; i++) {
        scanf("\n%lld %lld %lld", &oper, &from, &to);
        switch (oper) {
            case 1: // 1이면 숫자 교체
                from += base;
                int change = to - tree[from]; // 원래 있던 숫자를 빼고 새로 들어갈 숫자를 더한다.
                while(from != 0) {
                    tree[from] += change; // 현재 노드 값 변경
                    from /= 2; // 현재 노드의 부모 노드로
                }
                break;
            case 2: // 2이면 부분 합 구하기
                result = 0; // 합을 0으로 초기화
                from += base; // 제일 아래 위치한
                to += base;
                while (from <= to) {
                    if (from % 2 == 1) // 두개의 형제 노드 중 우측 노드면
                        result += tree[from++]; // 해당 값을 더하고 이후 노드로 이동(우측으로 한칸)
                    if (to % 2 == 0) // 두개의 형제 노드 중 좌측 노드면
                        result += tree[to--]; // 해당 값을 더하고 이전 노드로 이동(좌측으로 한칸)
                    from /= 2; // 현재 노드의 부모 노드로
                    to /= 2; // 현재 노드의 부모 노드로
                }
                printf("%lld\n", result);
                break;
        }
    }
    
    free(tree); // 동적 할당 해제
    return 0;
}