[Algorithm] Segment Tree (구간 트리)

Segment Tree (구간 트리)

구간 트리는 배열의 특정 구간에 대한 질의를 빠르게하기 위해 사용할 수 있다.

예를 들어 어떤 구간의 최소치를 구하고 싶다고 해보자. 보통 사용하는 방법은 구간 내 숫자를 하나하나 확인하면서 최소치를 찾는 것으로, O(N)의 시간이 걸린다. 이때 배열을 전처리해 구간 트리를 사용하면 O(N log N)의 시간에 찾을 수 있다.

구간 트리는 '주어진 배열의 구간들을 표현'하는 이진 트리를 만든다. 루트는 배열의 전체구간을 표현하고, 한 트리의 왼쪽 자식과 오른쪽 자식은 해당 구간의 왼쪽 반과 오른쪽 반을 표현한다. 각 노드는 해당 구간에 대한 계산 결과를 저장해둔다.

이렇게 전처리해둔 트리가 있다면 어떤 구간 [i, j]가 주어지면 이 구간을 구간 트리의 노드에 포함된 구간들의 합집합으로 표현할 수 있다. 예를 들어 [0, 10]으로 만든 구간 트리에서 [3, 7]은 [2], [3, 5], [6, 7] 로 표현할 수 있다.

 

프로세스

구간 트리는 위의 이미지와 같이 주어진 객체(배열)을 절반으로 나눠 결국 요소 하나만 남을 때까지 반복해서 자료로 저장한다. 만약 저 중 0~4번까지의 데이터가 필요하다면 2번 노드와 12번 노드에만 접근하면 된다. 모두 순회 할 필요가 없는 것이다.

 

특정 구간의 최소치를 찾는 구간 트리 구현

위에서 들었던 예시를 구현해보자.

먼저 트리를 저정할 배열을 만들어야 한다. 배열에는 모든 노드가 들어갈 수 있을 만큼 충분히 큰 크기여야 한다. 배일의 길이와 가장 가까운 2의 제곱으로 올림한 뒤 2를 곱하거나, 귀찮다면 그냥 4배 해주면 된다.

예시에서는 배열의 길이가 10이라고 가정한다.

이제 구간 트리를 만들어야 한다. 현재 노드에 저장한 값은 왼쪽과 오른쪽 자식 중 더 작은 것을 저장하면 된다.

const segmentTree = Array(10 * 4);

function init(start, end, node) {
    if (start === end) return segmentTree[node] = arr[start];
    
    const mid = Math.floor((start + end) / 2);
    return segmentTree[node] = Math.min(
        init(start, mid, node * 2), 
        init(mid + 1, end, node * 2 + 1)
    );
}

구간 트리를 초기화 과정은 O(N)의 시간 복잡도를 가진다.

임의의 구간에 대해 질의하기 위한 준비는 모두 마쳤다. 이제 질의를 해보자. 위에서 임의의 구간은 구간 트리 노드의 합집합으로 표현할 수 있다고 했다. 즉, start 와 end 가 범위 안에 있는 경우만 찾아주면 된다.

function query(left, right) {
    function helper(start, end, node, left, right) {
        // 범위 밖
        if (end < left || right < start) return Infinity;
        
        // 범위 안
        if (left <= start && end <= right) return segmentTree[node];
        
        // 그렇지 않으면 두 부분으로 나누어 구함
        const mid = Math.floor((start + end) / 2);
        return Math.min(
            helper(start, mid, node * 2, left, right),
            helper(mid + 1, end, node * 2 + 1, left, right)
        );
    }
    return helper(0, 9, 1, left, right);
}

트리에 대한 질의는 O(log N)의 시간 복잡도를 가지게 된다.

만약 배열의 요소가 수정되면 어떻게 할까? 이때는 해당 요소를 포함하고 노드들만 갱신해주면 된다.

function update(target, newValue) {
    function helper(start, end, node, target, newValue) {
        // 해당 구간과 상관 없으면 무시
        if (target < start || end < target) return segmentTree[node];
        
        // 리프 노드면 갱신
        if (start === end) return segmentTree[node] = newValue;
        
        const mid = Math.floor((start + end) / 2);
        return segmentTree[node] = Math.min(
            helper(start, mid, node * 2, target, newValue),
            helper(mid + 1, end, node * 2 + 1, target, newValue),
        )
    }
    return helper(0, 9, 1, target, newValue);
}

트리에 대한 수정 또한 O(log N)의 시간 복잡도를 가지게 된다.