자바 코딩테스트 강좌, 세그먼트 트리

안녕하세요! 이번 글에서는 자바를 사용한 코딩테스트에서 유용하게 쓰일 수 있는 알고리즘 중 하나인 세그먼트 트리에 대해 알아보겠습니다. 세그먼트 트리는 주로 구간 쿼리와 구간 업데이트를 효율적으로 처리하기 위해 사용됩니다. 이 글에서는 세그먼트 트리의 기본 개념과 이에 대한 예제 문제, 그리고 그 문제의 해결 과정을 자세히 설명드리겠습니다.

세그먼트 트리란?

세그먼트 트리는 배열에서 특정 구간의 정보를 저장하고 쿼리를 효율적으로 처리하기 위한 데이터 구조입니다. 주로 다음과 같은 작업을 지원합니다:

  • 지정한 구간의 합을 구하기
  • 지정한 위치의 값을 업데이트 하기
  • 구간의 최솟값, 최댓값을 구하기

세그먼트 트리는 O(log n) 시간에 쿼리를 처리할 수 있으며, 공간 복잡도는 O(n)입니다. 이를 통해 대규모의 데이터를 효과적으로 다룰 수 있습니다.

문제 소개

이제 ‘세그먼트 트리’를 활용한 문제를 하나 풀어보겠습니다.

문제 설명

주어진 정수 배열이 있을 때, 특정 구간의 합을 구하고, 특정 인덱스에 있는 값을 업데이트 하는 두 가지 기능을 제공하는 프로그램을 작성하세요.

구현해야 하는 기능은 다음과 같습니다:

  1. 구간 합 쿼리: 두 인덱스 lr가 주어지면, lr 사이의 합을 반환합니다.
  2. 업데이트 쿼리: 인덱스 index의 값을 value로 업데이트 합니다.

입력 형식

첫 번째 줄: 배열의 크기 n (1 ≤ n ≤ 10^5)
두 번째 줄: n개의 정수 (1 ≤ arr[i] ≤ 10^9)
세 번째 줄: 쿼리의 수 q
다음 q줄: 쿼리는 두 가지 형식 중 하나입니다:
    1. 1 l r (구간 합 쿼리)
    2. 2 index value (업데이트 쿼리)

출력 형식

구간 합 쿼리에 대한 결과를 한 줄씩 출력합니다.

문제 해결 과정

1. 세그먼트 트리 구현

먼저 세그먼트 트리를 구현해야 합니다. 세그먼트 트리는 비상위 노드로부터 리프 노드까지의 합을 저장하는 방식으로 구성됩니다. 이를 위한 배열 tree를 선언하고, 트리를 구축하는 메소드를 작성하겠습니다.

java
public class SegmentTree {
    private int[] tree;
    private int n;

    public SegmentTree(int[] arr) {
        n = arr.length;
        tree = new int[n * 4];
        buildTree(arr, 0, 0, n - 1);
    }

    private void buildTree(int[] arr, int node, int start, int end) {
        if (start == end) {
            tree[node] = arr[start];
        } else {
            int mid = (start + end) / 2;
            buildTree(arr, node * 2 + 1, start, mid);
            buildTree(arr, node * 2 + 2, mid + 1, end);
            tree[node] = tree[node * 2 + 1] + tree[node * 2 + 2];
        }
    }

    // 구간 합 쿼리 메소드
    public int rangeSum(int l, int r) {
        return rangeSumHelper(0, 0, n - 1, l, r);
    }

    private int rangeSumHelper(int node, int start, int end, int l, int r) {
        if (r < start || end < l) {
            return 0; // 쿼리 범위가 노드 범위와 겹치지 않음
        }
        if (l <= start && end <= r) {
            return tree[node]; // 노드 범위가 쿼리 범위에 포함됨
        }
        int mid = (start + end) / 2;
        int leftSum = rangeSumHelper(node * 2 + 1, start, mid, l, r);
        int rightSum = rangeSumHelper(node * 2 + 2, mid + 1, end, l, r);
        return leftSum + rightSum;
    }

    // 업데이트 메소드
    public void update(int index, int value) {
        updateHelper(0, 0, n - 1, index, value);
    }

    private void updateHelper(int node, int start, int end, int index, int value) {
        if (start == end) {
            tree[node] = value;
        } else {
            int mid = (start + end) / 2;
            if (start <= index && index <= mid) {
                updateHelper(node * 2 + 1, start, mid, index, value);
            } else {
                updateHelper(node * 2 + 2, mid + 1, end, index, value);
            }
            tree[node] = tree[node * 2 + 1] + tree[node * 2 + 2];
        }
    }
}

2. 메인 프로그램

세그먼트 트리 클래스를 포함하여 입력을 받아 쿼리를 처리하는 메인 프로그램을 작성하겠습니다.

java
import java.util.Scanner;

public class Main {
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        
        int n = scanner.nextInt();
        int[] arr = new int[n];
        
        for (int i = 0; i < n; i++) {
            arr[i] = scanner.nextInt();
        }

        SegmentTree segmentTree = new SegmentTree(arr);
        
        int q = scanner.nextInt();
        for (int i = 0; i < q; i++) {
            int type = scanner.nextInt();
            if (type == 1) { // 구간 합 쿼리
                int l = scanner.nextInt();
                int r = scanner.nextInt();
                System.out.println(segmentTree.rangeSum(l, r)); 
            } else if (type == 2) { // 업데이트 쿼리
                int index = scanner.nextInt();
                int value = scanner.nextInt();
                segmentTree.update(index, value);
            }
        }
        
        scanner.close();
    }
}

3. 시간 복잡도 분석

세그먼트 트리의 시간 복잡도는 다음과 같습니다:

  • 구간 합 쿼리: O(log n)
  • 업데이트 쿼리: O(log n)

따라서, n개의 데이터를 가진 배열에서 q개의 쿼리를 처리하면 총 O(q log n) 시간 복잡도를 가집니다.

맺음말

이번 강좌에서는 세그먼트 트리에 대해 자세히 알아보았습니다. 세그먼트 트리는 다양한 문제를 효율적으로 해결할 수 있는 강력한 도구이며, 카운팅 문제나 최댓값/최솟값 문제 등 여러 분야에서 활용될 수 있습니다. 이 자료가 여러분의 코딩테스트 준비에 도움이 되기를 바랍니다! 추가적인 질문이 있으시면 댓글로 남겨주세요.