세그먼트 트리(Segment Tree) 정의 및 구현 방법

세그먼트 트리(Segment Tree)의 필요성

특정 구간 내 쿼리를 효율적으로 처리하기 위한 자료구조다.

 

✅ 구간 쿼리

: 배열 또는 리스트와 같은 시퀀스 데이터에서 특정 구간에 대한 정보를 검색하거나 수정하는 작업을 말한다.

 

 

위와 같이 arr[50]의 배열이 있다고 가정했을 때, arr[3] ~ arr[10]까지의 합을 구하고, 특정 값을 변경하려는 행위가 필요하다고 가정해보자.

 

1. 단순 반복문을 통한 덧셈 연산

이 방법은 O(N)의 시간복잡도를 가지고 있으며, 이를 M번 수행한다고 가정한다면 O(NM)의 시간복잡도를 가지게 된다.

특정 값 A → B로 변경하는 것은 O(1)의 시간복잡도를 가지고, 이를 M번 수행한다면 O(M)의 시간복잡도를 가진다.

 

▶ 시간복잡도: O(NM) + O(M) = O(NM)

 

2. 누적합

이 방법은 최초에 누적합을 저장하기 위해 O(N)의 시간복잡도를 가지고, 이를 M번 수행한다면 O(N+M)의 시간복잡도를 가진다.

 

arr[3]의 데이터가 변경된다면 누적합 배열에도 문제가 생기기 때문에 변경이 이루어져야한다. M번의 데이터 변경이 생기면 O(NM)의 시간복잡도를 가진다.

 

▶ 시간복잡도: O(N+M) + O(NM) = O(NM)

 

두 가지 방법은 숫자가 커질수록 비효율적이며 짧은 시간이 주어진 알고리즘 문제를 해결할 수 없다.

하지만 세그먼트 트리를 이용하면 O(MlogN)의 시간복잡도로 시간 문제를 해결할 수 있다.

 

 

세그먼트 트리(Segment Tree) 

일반적으로 이진 트리(Binary Tree) 형태를 구성된다.

즉, 모든 노드는 1-2개의 자식(Child) 노드를 보유하고 있고, 부모(Parent) 노드는 자식 노드들의 합을 저장하고 있는 방식이다.

 

루트(Root) 노드의 번호는 1번이고, 현재 노드의 번호가 M이라면 왼쪽 자식 노드의 번호는 2M, 오른쪽 자식 노드의 번호는 2M+1이다.

 

크기가 5인 배열을 세그먼트 트리로 구현할 경우 아래와 같은 형태를 가진다.

 

위 그림은 세그먼트 트리에서 필요한 부분만 표현했다.

실제로 15번 노드까지 존재하며 10번부터 15번은 빈 값으로 입력된다.

 

트리에 필요한 노드 수는 2N -1이다.

2, 4, 8, 16, ··· 2의 제곱이 아니라면 불필요한 노드는 비워둔다. 

 

트리의 높이는 루트 노드에서 가장 긴 경로를 의미한다.

위 세그먼트 트리의 높이는 3이다.

 

arr배열의 크기가 5이므로 log5 → 2.xxxx. 2의 제곱이 아니기 때문에 +1을 해준다.

arr배열의 크기가 8이라면 log8 → 3. 2의 제곱이므로 3이다.

 

세그먼트 트리를 배열로 구현했을 때, 크기는 $2^{(h+1)}$ -1이다.

 

 

세그먼트(Segment Tree) 구현

 

✅ 생성 및 초기화

class SegmentTree {
    long tree[];
    int treeSize;
    
    public SegmentTree (int arrSize) {
        // 트리의 높이 구하기
        int h = (int) Math.ceil(Math.log(arrSize) / Math.log(2));
        
        this.treeSize = (int) Math.pow(2, h+1);
        tree = new long[treeSize];
    }
    
    public long init (long[] arr, int node, int start, int end) {
        // start == end라면 leaf 노드이므로 배열의 값 그대로 저장
        if (start == end) {
		    return tree[node] = arr[start];
        }
        
        // leaf 노드가 아니라면 자식 노드의 
        return tree[node] = init(arr, node*2, start, (start+end)/2) + init(arr, node*2+1, (start+end)/2+1, end);
    }
}

세그먼트 트리의 그림을 보면 알 수 있듯이 세그먼트 트리의 높이는 전달되는 배열의 크기에 영향을 받는다. 

Math.ceil 함수로 2의 제곱근이 아닌 경우에 +1을 해주기 위해서 사용해준다.

 

트리의 높이를 구해준 후에 재귀 형태로 리프 노드에 배열로 전달된 값을 입력하고, 해당 노드의 합을 부모 노드에 입력할 수 있다.

 

init(arr, 1, 1, arr.length-1); // init(원소배열, 루프노드(1), 원소배열 시작 idx, 원소배열 끝 idx);

위와 같이 init 메소드를 호출하게 되면 다음 노드의 인덱스 값(왼쪽 자식 노드(2*node)), 오른쪽 좌식 노드(2*node+1))을 찾아간다.

 

계속 해서 다음 노드를 찾아서 내려가다보면 start == end 의 경우가 생긴다.

리프 노드에 도달했다는 의미로 배열의 값을 반환하여 부모 노드의 값들을 하나씩 채워간다.

 

 

✅ 데이터 변경 update

public void update (int node, int start, int end, int idx, int diff) {
    // 변경할 idx 값이 범위를 벗어나면 확인할 필요 X
    if (idx < start || end < idx) return;
    
    // 변경된 차이만큼 각 노드에 저장
    tree[node] += diff;
    
    // 리프 노드가 아니라면 아래 자식 노드도 확인
    if (start != end) {
        update(node*2, start, (start+end)/2, idx, diff);
        update(node*2+1, (start+end)/2+1, end, idx, diff);
    }
}

중간값이 변경됐을 경우 재귀 형태로 기존의 값과 현재 값의 차이만큼 영향 받는 모든 노드에 더해서 저장해준다.

 

자식 노드의 합 범위가 현재 변경된 idx와 관련이 없다면 확인을 하지않으며, arr 배열의 데이터가 변경되어서 세그먼트 트리의 노드가 변경된 것이므로 arr 배열의 데이터도 변경해주어야한다.

 

 

✅ 구간 합 구하기

public long sum (int node, int start, int end, int left, int right) {
    // 범위를 벗어나는 경우 확인 X
    if (left > end || right < start) {
        return 0;
    }
    
    // 범위 내에 포함 시에는 더 내려가지 않음
    if (left <= start && end <= right) {
        return tree[node];
    }
      
    return sum(node*2, start, (start+end)/2, left, right) + sum(node*2+1, (start+end)/2+1, end, left, right);
}

구간 합을 구할 때는 4가지의 경우가 있다.

 

1. [left, right]와 [start, end]가 겹치지않는 경우 - 첫 번째 조건문

: 겹치지 않기 때문에 해당 구간은 불필요함 (범위를 벗어남)

 

2. [left, right]가 [start, end]와 완전히 일치하는 경우 - 두 번째 조건문

: 현재 그 구간합을 보유한 노드가 정답

 

3. [start, end]가 [left, right]를 완전히 포함하거나 한 쪽만 겹치는 경우 - 마지막 return

: left/right의 정확한 구간보다 더 넓은 범위이므로 자식 노드를 탐색한다.

 

 

 

<전체코드>

class SegmentTree {
    long tree[];
    int treeSize;
    
    public SegmentTree (int arrSize) {
        // 트리의 높이 구하기
        int h = (int) Math.ceil(Math.log(arrSize) / Math.log(2));
        
        this.treeSize = (int) Math.pow(2, h+1);
        tree = new long[treeSize];
    }
    
    public long init (long[] arr, int node, int start, int end) {
        // start == end라면 leaf 노드이므로 배열의 값 그대로 저장
        if (start == end) {
            return tree[node] = arr[start];
        }
        
        // leaf 노드가 아니라면 자식 노드의 
        return tree[node] = init(arr, node*2, start, (start+end)/2) + init(arr, node*2+1, (start+end)/2+1, end);
    }
    
    public void update (int node, int start, int end, int idx, int diff) {
        // 변경할 idx 값이 범위를 벗어나면 확인할 필요 X
        if (idx < start || end < idx) return;
    
        // 변경된 차이만큼 각 노드에 저장
        tree[node] += diff;
    
        // 리프 노드가 아니라면 아래 자식 노드도 확인
        if (start != end) {
            update(node*2, start, (start+end)/2, idx, diff);
            update(node*2+1, (start+end)/2+1, end, idx, diff);
        }
    
    public long sum (int node, int start, int end, int left, int right) {
        // 범위를 벗어나는 경우 확인 X
        if (left > end || right < start) {
            return 0;
        }
    
        // 범위 내에 포함 시에는 더 내려가지 않음
        if (left <= start && end <= right) {
            return tree[node];
        }
      
        return sum(node*2, start, (start+end)/2, left, right) + sum(node*2+1, (start+end)/2+1, end, left, right);
    }
}

 

 

 

 

 

 

 

---

참고: https://hongjw1938.tistory.com/20