[알고리즘] 세그먼트 트리(Segment Tree) 알고리즘 동작 원리 및 구현하기

세그먼트 트리 알고리즘

개요

• 세그먼트 트리 알고리즘은 트리 형태의 자료구조를 사용한다.
• 특정 구간 내 데이터에 대한 연산(쿼리)을 빠르게 구할 수 있는 트리 알고리즘이다.

키워드

• 구간
• 범위

시간복잡도

• 데이터 변경 : O(logN)
• 연산 : O(logN)
• 데이터 변경 시마다 M번 연산 : O((logN + logN) * M) = O(MlogN)

특징

누적 합을 담은 트리라고 가정하자.

N = 10 (노드의 개수 : 10) 인 그림에서 노드안에 적힌 순서는 값이 아닌 각 노드의 고유 번호라고 했을 때, 각 노드가 저장하고 있는 합의 범위(구간)를 나타낸 것이다.

각 노드는 다음과 같은 의미를 가진다.

• 리프 노드 : 해당 노드의 값 자체
• 리프 노드가 아닌 노드 : 왼쪽 자식과 오른쪽 자식의 합(누적 합 예시)

그 외 특징

• 리프 노드를 제외한 모든 노드는 항상 2개의 자식을 가진다. 
• 세그먼트 트리는 Full Binary Tree 형태를 가진다. 만약 N이 2의 제곱꼴인 경우는 Perfect Binary Tree가 된다.
• 리프 노드가 N개인 경우, 리프 노드가 아닌 노드가 N-1 개가 있다.

트리 종류	노드의 개수	높이
Full Binary Tree	2N-1	logN
Perfect Binary Tree	2N-1	2^(H+1) - 1

• (번외 꿀팁) 입력이 10만 이상인 경우 세그먼트 트리(혹은 그래프) 알고리즘 문제일 가능성이 크다.

구현 방법

문제

백준 알고리즘 (https://www.acmicpc.net/problem/2042)

• N개의 수가 있다.
• 중간에 수의 변경이 빈번히 일어나며, 그 중간 어떤 부분의 합을 구하려 한다.

다음의 예를 들었을 때 출력은 17이 된다.

• 1,2,3,4,5 라는 수가 있다
• 3번째 수를 6으로 바꾼다
• 2번째부터 5번째까지의 합을 구한다. (출력 : 17)

위와 같은 상태에서 다음과 같은 변경이 더 일어났을 경우, 출력은 12가 된다.

• 5번째 수를 2로 바꾼다
• 3번째부터 5번째까지의 합을 구한다. (출력 : 12)

세그먼트 트리 구현 흐름

• 생성 및 구성     : 생성자 + init
• 데이터 변경      : update 
• 구간 합 구하기 : sum

세그먼트 트리 구현

public class SegmentTree {

    private static long tree[];
    private static int treeSize;

    public SegmentTree(int size) {
        /** 배열의 크기인 tree size를 구하기 위한 높이 구하는 공식 */
        int height = (int) Math.ceil(Math.log(size)/Math.log(2));
        this.treeSize = (int) Math.pow(2, height+1);

        /** 공식을 외우기 어렵다면, input * 4로 해도 충분하다.. */
        this.treeSize = size*4;
        this.tree = new long[treeSize];
    }

    /**
     * 세그먼트 트리 값 초기화
     * @param 원소배열, 루프노드, 원소배열 시작 인덱스, 원소배열 끝 인덱스
     * @return 원소배열 값을 바탕으로 세그먼트 트리의 각 노드에 맞게 누적합을 세팅
     * */
    public long init(long[] arr, int node, int start, int end) {
        /** 배열의 시작과 끝이 같은 경우, leaf 노드이므로 원소 값 그대로 담기 */
        if(start == end) {
            return tree[node] = arr[start];
        }

        /** leaf 노드 아닌 경우, 자식노드의 합 담기 (재귀형태로 자식노드 세팅) */
        long leftChild = init(arr, node*2, start, (start+end)/2);
        long rightChild = init(arr, node*2+1, (start+end)/2+1, end);
        return tree[node] = leftChild + rightChild;
    }

    /**
     * 데이터 변경
     * @param 현재노드, 인덱스, 배열의 시작, 배열의 끝, 변경된 데이터의 인덱스, 기존 데이터와 변경 데이터 값의 차이
     * */
    public void update(int node, int start, int end, int index, long diff){

        /** 변경할 인덱스 값이 출력할 합의 범위 바깥인 경우 확인 불필요 */
        if(index < start || end < index) return;

        /** 변경된 값의 차이 저장 */
        tree[node] += diff;

        /** 현재 노드가 leaf 노드가 아닌 경우, 자식노드도 확인 */
        if(start != end) {
            update(node*2, start, (start+end)/2, index, diff);      // 왼쪽 자식노드
            update(node*2+1, (start+end)/2+1, end, index, diff);    // 오른쪽 자식노드
        }

        /** 이 때, 원소배열의 데이터 변경도 해놔야 같은 위치 인덱스 변경할 때 정확한 데이터값의 차이가 나온다 */
    }

    /**
     * 구간 합 구하기
     * @param 현재 노드, 배열의 시작, 배열의 끝, 누적합 범위의 시작 인덱스, 누적합 범위의 끝 인덱스스
     * */
   public long sum(int node, int start, int end, int left, int right) {

       /** 범위를 벗어나게 되는 경우 더할 필요 없음 */
       if(left > end || right < start){
           return 0;
       }

       /** 범위 내 완전히 포함 시에는 더 내려가지 않고 바로 리턴 */
       if(left <= start && end <= right){
           return tree[node];
       }

       /** 그 외의 경우 좌 / 우측으로 지속 탐색 수행 */
       long leftSum = sum(node*2, start, (start+end)/2, left, right);
       long rightSum = sum(node*2+1, (start+end)/2+1, end, left, right);
       return leftSum + rightSum;
    }
}

 

마지막 구간 합을 구하는 과정에서 구간의 경우의 수를 정리해보면,

노드에 저장된 구간이 [start - end] 이고, 합을 구해야하는 구간이 [left - right] 라고 할 때, 다음과 같은 4가지 경우로 나뉘어진다.

1. [left - right] 와 [start - end] 범위가 겹치지 않는 경우
2. [left - right] 가 [start - end] 를 완전히 포함하는 경우
3. [start - end] 가 [left - right] 를 완전히 포함하는 경우
4. [left - right] 와 [start - end] 가 겹쳐져 있는 경우