[백준] 2042 - 구간 합 구하기(세그먼트 트리)

2042 - 구간 합 구하기

1. 문제

문제
어떤 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(1 ≤ b ≤ N)번째 수를 c로 바꾸고 a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.

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

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

2. 요구사항

• 입력: 수의 개수 N(1<=N<=1000000), 변경 횟수 M, 구간 합 횟수 K
• 연산:
  1. 데이터 변경: 특정 인덱스($b$)의 수를 $c$로 바꾼다.
  2. 구간 합 구하기: 인덱스 $b$부터 $c$까지의 합을 구한다.
• 제약 사항: 입력되는 수는 $-2^{63}$ ~ $2^{63}-1$ 범위의 정수이다. (매우 중요: long 타입 필수)
• 목표: 빈번한 수의 변경과 구간 합 구하기를 시간 내에 처리해야 한다.

3. 접근 방식

왜 세그먼트 트리인가?

일반적인 1차원 배열을 사용하면 다음과 같은 시간 복잡도를 가진다.

• 수 변경: O(1)
• 구간 합 구하기: O(N) (최악의 경우 100만 번 연산)

하지만 문제에서 M과 K가 최대 1만 번씩 주어진다. 만약 구간 합을 구할 때마다 O(N)이 걸린다면 총 시간 복잡도는 O(K \times N)이 되어 시간 초과(Time Limit Exceeded)가 발생한다. 따라서 변경과 합 구하기를 모두 O(\log N)에 처리할 수 있는 세그먼트 트리를 사용해야 한다.

구현 설계트리
1. 크기: 배열로 트리를 구현할 때, N개의 데이터가 있다면 트리의 높이와 리프 노드 공간을 고려해 N \times 4 크기로 할당하는 것이 안전하다.
2. 자료형: 입력값의 범위가 long 범위이므로, 배열(arr, tree), 변수(diff, sum), 파싱(Long.parseLong) 등 관련된 모든 곳에 int가 아닌 long을 사용해야 한다.

핵심 함수

• init: 재귀적으로 왼쪽 자식과 오른쪽 자식의 합을 더해 현재 노드를 채운다.
• update: 변경된 값과 기존 값의 차이(diff)를 구해서, 해당 데이터가 포함된 모든 부모 노드에 diff를 더해준다. (리프 노드 도달 시 종료 조건 필수!)
• sum: 구하려는 범위(left, right)가 현재 노드 범위(start, end)에 완전히 포함되는지, 걸쳐있는지, 벗어나는지 판별하여 합을 반환한다.

3. 트러블 슈팅
   (시행착오)처음에는 코드가 계속 런타임 에러나 오답을 뱉었는데, 원인은 크게 두 가지였다.

(1) 자료형 문제 (NumberFormat Error)문제의 조건인 -2^{63} 범위를 간과하고 Integer.parseInt와 int[]를 사용했다. 입력값이 int 범위를 넘어서자마자 런타임 에러가 발생했다.-> 해결: 관련된 모든 변수와 배열을 long으로 변경했다.

(2) 재귀 종료 조건 누락 (IndexOutOfBounds)update 함수에서 리프 노드(start == end)에 도달했을 때 재귀를 멈추는 return; 문을 빼먹었다. 이로 인해 인덱스가 계속 커져서 4000000을 넘어버렸다.-> 해결: if (start == end) return; 조건을 추가하여 해결했다.

4. 코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

class SegmentTree {
    long[] tree;

    public SegmentTree() {
        tree = new long[4000000];
    }

    public long init(long[] arr, int node, int start, int end) {
        if (start == end) {
            return tree[node] = arr[start];
        }

        return tree[node] = init(arr, node * 2, start, (start + end) / 2)
                + init(arr, node * 2 + 1, (start + end) / 2 + 1, end);
    }
}

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;
        st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        int M = Integer.parseInt(st.nextToken());
        int K = Integer.parseInt(st.nextToken());

        long[] arr = new long[N + 1];
        for (int i = 1; i <= N; i++) {
            st = new StringTokenizer(br.readLine());
            arr[i] = Long.parseLong(st.nextToken());
        }

        SegmentTree segmentTree = new SegmentTree();
        segmentTree.init(arr, 1, 1, N);

        for (int i = 0; i < K + M; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            long c = Long.parseLong(st.nextToken());

            if(a == 1) {
                long diff = c - arr[b];
                arr[b] = c;
                update(segmentTree.tree, 1, 1, N, b, diff);
            }
            else {
                long answer = sum(segmentTree.tree, 1, 1, N, b, (int)c);
                System.out.println(answer);
            }
        }
    }

    static void update(long[] tree, int node, int start, int end, int idx, long diff) {
        if (idx < start || idx > end) {
            return;
        }

        tree[node] += diff;

        if(start == end) return;

        update(tree, node * 2, start, (start + end) / 2, idx, diff);
        update(tree, node * 2 + 1, (start + end) / 2 + 1, end, idx, diff);
    }

    static long sum(long[] tree, int node, int start, int end, int left, int right) {
        if(left > end || right < start) {
            return 0;
        }

        if(left <= start && right >= end) {
            return tree[node];
        }

        return sum(tree, node * 2, start, (start + end) / 2, left, right)
                + sum(tree, node * 2 + 1, (start + end) / 2 + 1, end, left, right);
    }

}

5. 느낀점

"자료형 확인은 선택이 아니라 필수다."

단순히 알고리즘 로직만 맞다고 정답이 아니었다. 문제에 주어진 수의 범위를 보고 int를 쓸지 long을 쓸지 판단하는 것 또한 실력이라는 것을 뼈저리게 느꼈다.

또한 세그먼트 트리는 개념적으로는 '구간을 반으로 쪼갠다'는 것이 명확하지만, 실제 코드로 옮길 때 재귀 함수의 종료 조건(Base Case)이나 배열 인덱스 처리(node * 2) 등을 아주 꼼꼼하게 작성해야 한다는 것을 배웠다. 이번 문제를 통해 세그먼트 트리의 기본 템플릿을 확실히 익혔으니, 다음에는 응용 문제인 BOJ 1275번 등에 도전해봐야겠다.