세그먼트 트리(Segment Tree)

1. 개념 🛫

  세그먼트 트리(Segment Tree)란 여러 개의 데이터가 존재할 때 특정 구간의 합을 구하는 데 사용하는 자료구조이다.

  트리 종류 중 하나로 이진트리의 형태이며, 특정 구간의 합을 가장 빠르게 구할 수 있다는 장점이 있다. O(log N)

 

2. 동작방식 🚁

  세그먼트 트리를 이용하기 위해서는 우선 '구간 합 트리'를 생성해야 한다. 그리고 '구간 합을 구하는 함수'와 '원소를 수정하는 함수'를 이용하여 구간 합을 구하고 특정 원소를 수정할 수 있다. 

  '1 9 3 8 4 5 5 9 10 3 4 5'라는 12개의 원소를 가진 배열이 주어졌을 때 세그먼트 트리를 이용하여 구간 합을 구해보자.

1. 구간 합 트리 생성하기

  '구간'은 원래 데이터 범위를 반씩 분할하는 것으로 설정한다. 최상단 노드는 전체 원소를 포함하는 구간인 0~11로 설정된다. 최상단 노드의 왼쪽 자식은 0~5, 오른쪽 자식은 6~11의 구간을 가진다.

  이를 그림으로 표현하면 다음과 같은 구간을 가지게 된다. (여기서는 노드의 값을 표현한 게 아니라 구간을 표시한 것이다.)

  구간 합 트리의 원소 개수는 32개여야 한다. 이를 간단히 알 수 있는 방법은 데이터의 개수가 N(=12) 개 일 때 N보다 큰 제곱수(= 16)를 2배 한 것(=32)이다. 더 간단하게는 데이터의 개수 N(=12)에 4를 곱한 값(=48)을 배열의 크기로 할당해 놓으면 된다.

  각각의 노드 값들은 노드의 구간에 속한 원소들의 합을 나타낸다.

  이를 구현하는 방법은 재귀를 이용할 수 있다.

static int init(int startIdx, int endIdx, int nodeIdx) {
    // 자기 자신만을 원소로 가지는 구간 // 끝 구간
    if(startIdx == endIdx)
        return segTree[nodeIdx] = arr[startIdx];

    int mid = (startIdx + endIdx) / 2;
    // 왼쪽 자식 노드 구간 합
    int leftSum = init(startIdx, mid, nodeIdx * 2);
    // 오른쪽 자식 노드 구간 합
    int rightSum = init(mid + 1, endIdx, nodeIdx * 2 + 1);

    return segTree[nodeIdx] = leftSum + rightSum;

}

2. 구간 합을 구하는 함수 만들기

  4~8의 구간 합을 구한다고 하자. 그러면 아래의 세 노드의 값만 더해주면 된다.

  해당 노드들을 얻기 위해서는 3가지 경우가 존재한다. 찾고자 하는 구간을 타겟 구간이라고 하자.

1. 탐색하는 구간이 타겟 구간에 포함되는 경우
   해당 노드 값을 가져온다.
2. 탐색하는 구간이 타겟 구간에 포함되지 않는 경우
   더 이상 탐색하지 않는다.
3. 탐색하는 구간이 타겟 구간에 일부분 포함되는 경우
   자식 노드들의 구간을 다시 살펴보면서 1번과 2번의 경우에 해당할 때까지 탐색한다.

  이를 코드로 표현하면 아래와 같다.

// 구하고자 하는 구간 : left ~ right
static int sum(int startIdx, int endIdx, int nodeIdx, int left, int right) {
    // 구간을 벗어나는 경우
    if(left > endIdx || right < startIdx)
        return 0;

    // 구간에 포함되는 경우
    if(left <= startIdx && endIdx <= right)
        return segTree[nodeIdx];

    // 구간에 걸치는 경우
    // 자식 노드로 진입하여 구간에 속하는 것만 찾기
    int mid = (startIdx + endIdx) / 2;
    int leftRangeSum = sum(startIdx, mid, nodeIdx * 2, left, right);
    int rightRangeSum = sum(mid + 1, endIdx, nodeIdx * 2 + 1, left, right);
    return leftRangeSum + rightRangeSum;
}

3. 특정 원소의 값을 수정하는 함수 만들기

  특정 원소의 값을 수정한다고 하면 해당 원소를 포함하고 있는 구간의 노드들을 모두 업데이트해 주면 된다. 예를 들어 3번 인덱스의 값을 수정한다고 하면 아래 표시된 노드들을 모두 업데이트해줘야 한다. (동그라미 안에는 구간의 범위를 표시한 것이다.)

  이를 코드로 표현하면 아래와 같다.

static void update(int startIdx, int endIdx, int nodeIdx, int targetIdx, int diff) {
    // 구간을 벗어나는 경우
    if(targetIdx < startIdx || targetIdx > endIdx)
        return;

    // 구간에 속하는 경우(diff에 영향을 받는 경우) 
    // diff(기존값과 새로운 값의 차이)를 더해주기
    segTree[nodeIdx] += diff;

    // leaf 노드인 경우
    if(startIdx == endIdx)
        return;

    // 수정하고자 하는 노드를 찾지 못했기 때문에
    // targetIdx를 찾기위해 자식노드로 내려가야 함
    int mid = (startIdx + endIdx) / 2;

    update(startIdx, mid, nodeIdx * 2, targetIdx, diff);
    update(mid + 1, endIdx, nodeIdx * 2 + 1, targetIdx, diff);
}

 

3. 구현(JAVA) 🚀

static int[] arr = {1, 9, 3, 8, 4, 5, 5, 9, 10, 3, 4, 5}; 
static int cnt = arr.length;
static int[] segTree = new int[4 * cnt];

public static void main(String[] args) {

    init(0, cnt - 1, 1);

    // 0부터 12까지의 구간 합(전체 구간)
    System.out.println(sum(0, cnt - 1, 1, 0, 12)); // 66
    // 4부터 8까지의 구간 합(전체 구간)
    System.out.println(sum(0, cnt - 1, 1, 4, 8)); // 33

    // 3번 인덱스의 값을 8에서 7로 수정 (diff = -6)
    update(0, cnt - 1, 1, 3, -6);

    // 변경 후 : 0~2 구간합
    // 변화X
    System.out.println(sum(0, cnt - 1, 1, 0, 2)); // 13

    // 변경 후 : 전체 구간합
    // 변화O (66 -> 60)
    System.out.println(sum(0, cnt - 1, 1, 0, 12)); // 60
}

// 1. 구간 합 트리 생성하기
static int init(int startIdx, int endIdx, int nodeIdx) {
    // 자기 자신만을 원소로 가지는 구간 // 끝 구간
    if(startIdx == endIdx)
        return segTree[nodeIdx] = arr[startIdx];

    int mid = (startIdx + endIdx) / 2;
    // 왼쪽 자식 노드 구간 합
    int leftSum = init(startIdx, mid, nodeIdx * 2);
    // 오른쪽 자식 노드 구간 합
    int rightSum = init(mid + 1, endIdx, nodeIdx * 2 + 1);

    return segTree[nodeIdx] = leftSum + rightSum;

}

// 2. 구간 합을 구하는 함수 만들기
// 구하고자 하는 구간 : left ~ right
static int sum(int startIdx, int endIdx, int nodeIdx, int left, int right) {
    // 구간을 벗어나는 경우
    if(left > endIdx || right < startIdx)
        return 0;

    // 구간에 포함되는 경우
    if(left <= startIdx && endIdx <= right)
        return segTree[nodeIdx];

    // 구간에 걸치는 경우
    // 자식 노드로 진입하여 구간에 속하는 것만 찾기
    int mid = (startIdx + endIdx) / 2;
    int leftRangeSum = sum(startIdx, mid, nodeIdx * 2, left, right);
    int rightRangeSum = sum(mid + 1, endIdx, nodeIdx * 2 + 1, left, right);
    return leftRangeSum + rightRangeSum;
}

// 3. 특정 원소의 값을 수정하는 함수 만들기
static void update(int startIdx, int endIdx, int nodeIdx, int targetIdx, int diff) {
    // 구간을 벗어나는 경우
    if(targetIdx < startIdx || targetIdx > endIdx)
        return;

    // 구간에 속하는 경우(diff에 영향을 받는 경우) 
    // diff(기존값과 새로운 값의 차이)를 더해주기
    segTree[nodeIdx] += diff;

    // leaf 노드인 경우
    if(startIdx == endIdx)
        return;

    // 수정하고자 하는 노드를 찾지 못했기 때문에
    // targetIdx를 찾기위해 자식노드로 내려가야 함
    int mid = (startIdx + endIdx) / 2;

    update(startIdx, mid, nodeIdx * 2, targetIdx, diff);
    update(mid + 1, endIdx, nodeIdx * 2 + 1, targetIdx, diff);
}

 

4. 시간복잡도 🛸

  구간 합 트리를 생성하는 데 걸리는 시간은 O(N log N)이다.

  그리고 구간 합을 구하거나 값을 경신하는 데 걸리는 시간은 트리 구조의 특성을 이용하기 때문에 O(log N)이다.

 

해당 글은
나동빈 님의 '41. 세그먼트 트리(Segment Tree)',
pseong 님의 '세그먼트 트리(구간 트리), 느리게 갱신되는 세그먼트 트리(Lazy Propagation)'
를 참고하였습니다.