[백준 코딩테스트] 2042번 구간 합 구하기

풀이

처음에는 단순 반복문을 사용해 문제를 접근하려고 하니 시간복잡도로 인해 런타임에러가 발생하였다.
이 문제는 세그먼트 트리 알고리즘을 사용해 풀어야 하는 문제였다. 세그먼트 트리를 사용하면 시간복잡도가 O(logN) 이기 때문에 내가 마주한 시간초과 문제를 해결할 수 있었다.

세그먼트 트리

세그먼트 트리란 완전 이진 트리로, 모든 노드가 자식을 2개씩 가지며 마지막 레벨에선 왼쪽부터 차례대로 채워지는 트리를 말한다. n개의 노드를 가질 때 마지막 레벨은 최대 n 개의 리프 노드를 가질 수 있으며 완전 이진 트리의 총 노드수는 2*n - 1 개가 필요하다.

주어진 예제 입력을 통한 세그먼트 트리 초기화

init() 메소드를 사용해 초기화하는 과정이며 이를 트리구조로 나타내면 다음과 같은 구조를 확인할 수 있다.

diff를 구하는 이유 ?

차이값을 구해 세그먼트 트리에서 특정 값이 업데이트 되는 경우 해당 값이 속한 구간의 합을 효율적으로 갱신하기 위해서이다.
이를 통해 트리 일부만 수정하여 전체 구간의 합을 유지할 수 있다.

내 코드

import java.util.*;
import java.io.*;

public class Main {

    static long[] arr; // 원본 배열
    static long[] tree; // 세그먼트 트리 배열

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        ArrayList<Long> answer = new ArrayList<>();

        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());

        arr = new long[n + 1]; // 입력된 배열 (1-based index)
        tree = new long[4 * n]; // 세그먼트 트리 배열

        // 입력 값 저장
        for (int i = 1; i <= n; i++) {
            arr[i] = Long.parseLong(br.readLine());
        }

        // 세그먼트 트리 생성
        init(1, n, 1);

        // m + k개의 쿼리 처리
        for (int i = 0; i < m + k; i++) {
            st = new StringTokenizer(br.readLine());
            int type = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            long c = Long.parseLong(st.nextToken());

            if (type == 1) { // 값 변경
                long diff = c - arr[b];
                arr[b] = c;
                update(1, n, 1, b, diff);
            } else if (type == 2) { // 구간 합 계산
                answer.add(sum(1, n, 1, b, (int) c));
            }
        }

        for(long i : answer){
            System.out.println(i);
        }
    }

    // 세그먼트 트리 초기화
    static long init(int start, int end, int node) {
        if (start == end) {
            return tree[node] = arr[start];
        }

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

    // 특정 범위의 합 구하기
    static long sum(int start, int end, int node, int left, int right) {
        if (left > end || right < start) {
            return 0;
        }

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

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

    // 세그먼트 트리 업데이트
    static void update(int start, int end, int node, int idx, long diff) {
        if (idx < start || idx > end) {
            return;
        }

        tree[node] += diff;

        if (start != end) {
            int mid = (start + end) / 2;
            update(start, mid, node * 2, idx, diff);
            update(mid + 1, end, node * 2 + 1, idx, diff);
        }
    }
}