[BOJ] 2042. 구간 합 구하기

 

구간 합 구하는 방법은 합 배열의 개념을 이용해서 해결할 수 있으나 이번 문제와 같은 경우, 중간 중간 수가 자주 변경되는 상황이 일어나기 때문에 합 배열로 풀기에는 시간이 오래 걸린다는 단점이 있다. 그렇기에 이와 같은 문제는 세그먼트 트리를 이용해 풀이하는 것이 권장된다.

 

(세그먼트 트리의 경우, 개념 이해하는게 꽤 오래 걸렸으나 그래도 문제와 예제를 통해 세그먼트 트리 문제임을 알아채기가 편하다고 느껴졌다(?). 풀이 공식같은 것도 딱 정해져 있어서 차라리 그리디, 구현같은 것보다는 상대적으로 낫게 느껴진다..)

 

 

문제

 

 

 

 

예제

 

 

접근

 

 

트리 초기화하기

 

리프 노드의 개수가 데이터의 개수(N) 이상이 되도록 트리 배열을 만든다.

트리 배열의 크기 구하는 방법 : 2^k >= N 을 만족하는 k의 최솟값을 구한 후 2^k * 2를 트리 배열의 크기로 정의한다.

 

위 예제와 같이 수의 개수가 N=5일 때, k의 최솟값과 트리 배열의 크기 및 리프 노드의 시작 위치를 구하는 방법이다.

 

 

0부터 시작해서 15까지 트리 배열의 크기만큼 트리 값을 초기화 한다. 리프노드의 시작위치 - 1 인덱스부터 1번 방향으로 트리 배열의 값을 채워나가도록 한다. 7번부터 1번까지 노드의 값을 채우는데, 자식 노드의 인덱스는 이진 트리 형식이므로 2n, 2n+1이 된다. (위의 입력값인 N과 혼동 X)

 

구간합 노드 채우는 식 : A[n] = A[2n] + A[2n+1]

 

 

 

 

트리 갱신

 

a가 1일 때, b번째 수를 c로 바꾸고, a가 2일 때, b번째 수에서 c번째 수까지의 합을 구하여 출력해야 한다.

 

먼저 a가 1일 때, 갱신하는 로직부터 그림으로 표현해보았다.

 

 

 

 

 

start index로부터 3번째 수를 6으로 갱신하고, 그에 따라 부모 노드의 값이 갱신되는 것을 트리로 표현해보았다.

맨 아래 리프노드의 값에서 +3하여 6으로 변경되는 시점부터 시작한다.

 

 

 

 

 

구간합 구하기

 

다음은 a가 2일 때, 2~5까지의 구간합을 구하는 로직이다. 이를 위해서는 아래와 같은 질의값 구하는 방법이 적용된다.

 

1. start_idx % 2 == 1 일 때, 해당 노드 선택
2. end_idx % 2 == 0 일 때, 해당 노드 선택
3. start_idx depth 변경 => start_idx = (start_idx + 1) / 2 
4. end_idx depth 변경 => end_idx = (end_idx - 1) / 2
5. 위의 1~4번까지 반복하다 end_idx < start_idx가 되면 반복문 종료

 

위 과정을 적용하며 연산이 반복될 때마다 start 인덱스를 증가시키고, end 인덱스를 감소시키며 구간합을 계속 구해나간다.

그리고 두 개의 인덱스가 같아지는 순간 반복문을 종료하고 구간합 결과를 반환한다.

 

 

 

 

 

풀이

 

import java.util.Scanner;

public class BOJ2042 {

    // 세그먼트 트리 배열
    static long[] tree;
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        // 수 개수
        int N = sc.nextInt();
        // 수의 변경이 일어나는 횟수
        int M = sc.nextInt();
        // 구간의 합을 구하는 횟수
        int K = sc.nextInt();

        int treeHeight = 0;
        int length = N;

        while (length != 0) {
            length /= 2;
            treeHeight++;
        }

        // 트리 배열 크기 : 2^k * 2
        int treeSize = (int) Math.pow(2, treeHeight+1);
        // 2^k - 1 부터 1번 쪽으로 채움 (리프 노드 시작 인덱스)
        int leftNodeStartIndex = treeSize / 2 - 1;
        tree = new long[treeSize + 1];

        // 데이터 리프노드에 입력 (리프 노드 시작 위치 : 2^k)
        for (int i=leftNodeStartIndex+1; i<=leftNodeStartIndex+N; i++) {
            tree[i] = sc.nextLong();
        }

        // 초기 트리 생성
        setTree(treeSize-1);

        for (int i=0; i<M+K; i++) {
            long a = sc.nextLong();
            // 시작/끝 인덱스
            int start = sc.nextInt();
            long end = sc.nextLong();

            if (a == 1) {
                // 인덱스 값 바꾸기
                changeVal(leftNodeStartIndex + start, end);
            } else if (a == 2) {
                // start부터 end까지의 구간합 구함
                start = start + leftNodeStartIndex;
                end = end + leftNodeStartIndex;
                System.out.println(getSum(start, (int) end));
            } else {
                return;
            }
        }
    }

    private static long getSum(int start, int end) {
        // 부분합
        long partSum = 0;
        // 시작 인덱스와 종료 인덱스가 교차될 때까지
        while (start <= end) {
            // 해당 노드의 값 구간합에 추가하거나 시작 인덱스 증가
            if (start % 2 == 1) {
                partSum = partSum + tree[start];
                start++;
            }
            // 해당 노드의 값 구간합에 추가하거나 종료 인덱스 감소
            if (end % 2 == 0) {
                partSum = partSum + tree[end];
                end--;
            }

            start = start / 2;
            end = end / 2;
        }
        return partSum;
    }

    // index의 값을 변경하는 함수
    private static void changeVal(int index, long value) {
        long diff = value - tree[index];
        while (index>0) {
            tree[index] = tree[index] + diff;
            index = index/2;
        }
    }

    // 트리 노드 채움 (2^k - 1부터 1번쪽으로)
    private static void setTree(int i) {
        while (i != 1) {
            tree[i/2] += tree[i];
            i--;
        }
    }
}