세그먼트 트리 뽀개기

세그먼트 트리란? (Segment Tree)

여러개의 데이터가 연속적으로 존재할 때 특정한 범위에 대한 질의를 O(logN)의 시간복잡도로 해결할 수 있습니다.

여기서 말하는 질의(query)란 특정 index에 대한 update, 특정 구간에 대한 SUM, MIN, MAX등이 해당됩니다.

 

기본적인 개념은 데이터의 범위를 반씩 분할하며 그 구간의 query를 수행합니다.

 

구간합을 구하는 Segment Tree

위 그림에서와 같이 [1, 3, 5, 7, 9, 11]이라는 배열이 주어질 때, 구간의 대표값을 Tree에 저장합니다.

 

Root노드는 0 ~ 5까지의 index를 대표하는 값,

left child는 0 ~ 2 까지의 index를 대표하는 값,

right child는 3~5 까지의 index를 대표하는 값으로 저장합니다.

 

세그먼트 트리는 재귀로 구현하면 간단하게 구현할 수 있습니다.

 

 

세그먼트 트리에서 query를 조회할 때에는 해당 query의 범위가 해당 노드가 대표하고 있는 범위를 포함하는 지,

포함하지 않는 지의 여부가 중요합니다.

 

해당 노드가 대표하고 있는 범위가 query의 내부에 있다면 그 값을 선택해서 비교하고, 만약 노드가 대표하고 있는 범위가 query의 범위 밖에 있다면 해당 노드는 선택하지 않습니다.

 

주의할 점으로는

1. 처음에 min, max 등의 배열에 값을 할당하는 초기화 코드가 필요한데, 이 또한 재귀로 구현할 수 있습니다.

2. tree의 size는 주어진 데이터의 개수 n보다 큰 가장 가까운 2의 제곱수의 2배가 필요합니다.

따라서 보통은 넉넉하게 n * 4의 size로 할당합니다.

 

재귀에 대한 코드가 반복적이므로 한번 구현하면 편하게 사용할 수 있습니다.

public class SegmentTree {

    static int n;
    static int[] a, maxTree, minTree;

    // 세그먼트 트리 초기화
    //node: 해당노드
    //nodeLeft, nodeRight: 실제 배열에서 몇번째 인덱스부터 몇번째 인덱스까지를 대표하느냐?
    public static void init(int node, int nodeLeft, int nodeRight) {
        if (nodeLeft == nodeRight) {
            maxTree[node] = a[nodeLeft];
            minTree[node] = a[nodeLeft];
            return;
        }

        int mid = (nodeLeft + nodeRight) / 2;

        init(
            node * 2,
            nodeLeft,
            mid
        );

        init(
            node * 2 + 1,
            mid + 1,
            nodeRight
        );

        maxTree[node] = Math.max(maxTree[node * 2], maxTree[node * 2 + 1]);
        minTree[node] = Math.min(minTree[node * 2], minTree[node * 2 + 1]);
    }

    // 세그먼트 트리 갱신
    public static void update(int node, int nodeLeft, int nodeRight, int queryIndex, int value) {
        if (queryIndex < nodeLeft || nodeRight < queryIndex) {
            return;
        }

        if (nodeLeft == nodeRight) {
            maxTree[node] = value;
            minTree[node] = value;
            return;
        }

        int mid = (nodeLeft + nodeRight) / 2;

        update(
            node * 2,
            nodeLeft,
            mid,
            queryIndex,
            value
        );

        update(
            node * 2 + 1,
            mid + 1,
            nodeRight,
            queryIndex,
            value
        );

        maxTree[node] = Math.max(maxTree[node * 2], maxTree[node * 2 + 1]);
        minTree[node] = Math.min(minTree[node * 2], minTree[node * 2 + 1]);
    }

    // 최대값 쿼리
    public static int queryMax(int node, int nodeLeft, int nodeRight, int queryLeft, int queryRight) {
        if (queryRight < nodeLeft || nodeRight < queryLeft) {
            return 0;
        }

        if (queryLeft <= nodeLeft && nodeRight <= queryRight) {
            return maxTree[node];
        }

        int mid = (nodeLeft + nodeRight) / 2;
        int leftMax = queryMax(
            node * 2,
            nodeLeft,
            mid,
            queryLeft,
            queryRight
        );
        int rightMax = queryMax(
            node * 2 + 1,
            mid + 1,
            nodeRight,
            queryLeft,
            queryRight
        );

        return Math.max(leftMax, rightMax);
    }

    // 최소값 쿼리
    static int queryMin(int node, int nodeLeft, int nodeRight, int queryLeft, int queryRight) {
        if (queryLeft > nodeRight || queryRight < nodeLeft) {
            return Integer.MAX_VALUE;
        }

        if (queryLeft <= nodeLeft && nodeRight <= queryRight) {
            return minTree[node];
        }

        int mid = (nodeLeft + nodeRight) / 2;
        int leftMin = queryMin(
            node * 2,
            nodeLeft,
            mid,
            queryLeft,
            queryRight
        );
        int rightMin = queryMin(
            node * 2 + 1,
            mid + 1,
            nodeRight,
            queryLeft,
            queryRight
        );

        return Math.min(leftMin, rightMin);
    }

}