세그먼트 트리는 구간합이나 구간 최소원소를 찾는데 유리한 이진트리이다. 기본적으로 그룹을 이분탐색의 개념으로 나누고 구간에 대한 연산을 적용한 자료구조라고 볼 수 있다.

A = {1, 2, 3, ..., N}이라는 배열이 존재하며 두 가지 쿼리를 제공하는 문제가 있다.

1. i부터 j까지 범위의 합 구하기
2. i번째 배열의 값을 변경하는 update

이를 배열로 해결한다면 1. O(N), 2. O(1)이 된다. 두 연산을 M번 수행하면 총 O(MN)의 시간복잡도가 든다.

이를 세그먼트 트리로 해결한다면 두 연산 모두 O(logN)이 되며 총 O(MlogN)으로 해결 가능하다.

주의할 점은 공간복잡도가 약 O(4N)이기에 메모리를 많이 먹는다. -> 이를 개선한 다이나믹 세그먼트 트리라는 개념도 있다.

구현

이진트리를 배열로 구현하게 되면 부모 노드의 인덱스가 i라고 할때 왼쪽자식 노드는 2i이며 오른쪽 노드 자식은 2i + 1이다.

이를 활용해 배열 A가 있을때 세그먼트 트리는 아래 배열처럼 구현된다.

세그먼트 트리는 Complete Binary Tree이기 때문에 Tree 배열에 필요한 범위는 h = ceil(log_2 n)이라고 할때 Full Binary Tree기준 노드개수는 2^(h + 1) - 1이다.

S(n)은 세그먼트 트리에 필요한 노드 수이다.

편의를 위해 2^(h + 1) - 1이나 조금 넉넉하게 4n으로 크기를 정해준다.

    public SegmentTree(int[] A) {
        int h = (int)Math.ceil(Math.log(A.length));
        // 1. fit한 방법
        int nodeCnt = 1 << (h + 1);
        // 2. 넉넉 간단 방법
        nodeCnt = 4 * A.length;
        tree = new int[nodeCnt];
    }

연산

1. Build : 세그먼트 트리 초기화
2. Query : 구간 정보를 조회
3. Update : 배열의 값 변경

시간복잡도
1. O(2^(h + 1) - 1) < O(4n)
2. O(log n)
3. O(log n)

매개변수 설명

각 연산에서 매개변수 start, end, node는 그림에서 상태를 의미한다.

Build

재귀적으로 트리를 구현 가능

• 리프 노드라면 배열의 요소를 저장
• 내부 노드라면 구간 정보 저장

    public void build(int node, int start, int end) {
    	// 리프 노드라면 배열의 요소를 저장
        if (start == end) {
            tree[node] = A[start];
            return;
        }
		// 내부 노드라면 구간 정보 저장
        int mid = (start + end) / 2;
        build(node * 2, start, mid);
        build(node * 2 + 1, mid + 1, end);
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

매개변수명	설명
node	연산을 진행할 root node index
start	연산을 진행할 A배열 시작 인덱스
end	연산을 진행할 A배열 끝 인덱스

Query

• 노드가 지정된 범위 밖에 있다면 0 반환
• 노드가 나타내는 범위가 지정된 범위 내에 있다면 값 반환
• 노드가 나타내는 범위가 지정된 범위 일부만 포함한다면 왼쪽 자식 + 오른쪽 자식의 합 반환
  
  2에서 4까지의 구간합을 구한다면 그림같이 컬러 노드들을 방문하고, 그 중 초록색 노드의 값을 사용하여 결과를 반환.

    public int query(int node, int start, int end, int left, int right) {
    	// 노드가 지정된 범위 밖에 있다면 0 반환
        if (right < start || end < left) {
            return 0;
        }
        // 노드가 나타내는 범위가 지정된 범위 내에 있다면 값 반환
        if (left <= start && end <= right) {
            return tree[node];
        }
        
        // 노드가 나타내는 범위가 지정된 범위 일부만 포함한다면 왼쪽 자식 + 오른쪽 자식의 합 반환
        int mid = (start + end) / 2;
        int left_child = query(node * 2, start, mid, left, right);
        int right_child = query(node * 2 + 1, mid + 1, end, left, right);
        
        return left_child + right_child;
    }

매개변수명	설명
node	연산을 진행할 root node index
start	연산을 진행할 A배열 시작 인덱스
end	연산을 진행할 A배열 끝 인덱스
left	원하는 A배열 구간의 시작 인덱스
right	원하는 A배열 구간의 끝 인덱스

Update

배열 2의 값을 5로 변경하는 과정

매개변수명	설명
node	연산을 진행할 root node index
start	연산을 진행할 A배열 시작 인덱스
end	연산을 진행할 A배열 끝 인덱스
index	값을 변경할 A배열 기준 인덱스
val/diff	변경할 값/기존 값과의 차

구현1

• 리프노드라면 배열의 값을 변경
• 내부노드라면 구간 정보를 변경

    public void update(int node, int start, int end, int index, int val) {
    	// 리프노드라면 배열의 값을 변경
        if (start == end) {
            tree[node] = val;
            return;
        }
        // 내부 노드라면 구간 정보를 변경
        int mid = (start + end) / 2;
        if (start <= index && index <= mid) {
            update(node * 2, start, mid, index, val);
        } else {
            update(node * 2 + 1, mid + 1, end, index, val);
        }
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

구현2

• 탐색범위를 벗어나면 종료
• 변화량을 경로에 있는 노드에 모두 적용
• 리프노드가 아니라면 탐색 진행

    public void update(int node, int start, int end, int index, int diff) {
    	// 탐색범위를 벗어나면 종료
        if (index < start || end < index) {
            return;
        }
        
        // 변화량을 경로에 있는 노드에 모두 적용
        tree[node] += diff;
        
        // 리프노드가 아니라면 탐색 진행
        if (start != end) {
            int mid = (start + end) / 2;
            update(node * 2, start, mid, index, diff);
            update(node * 2 + 1, mid + 1, end, index, diff);
        }
    }

전체 코드

class SegmentTree {
    int[] tree;
    int[] A;

    public SegmentTree(int[] A) {
        this.A = A;
        int h = (int)Math.ceil(Math.log(A.length));
        int nodeCnt = 4 * A.length;
        tree = new int[nodeCnt];
    }

    public void build(int node, int start, int end) {
        if (start == end) {
            tree[node] = A[start];
            return;
        }

        int mid = (start + end) / 2;
        build(node * 2, start, mid);
        build(node * 2 + 1, mid + 1, end);
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

    public int query(int node, int start, int end, int left, int right) {
        if (right < start || end < left) {
            return 0;
        }
        if (left <= start && end <= right) {
            return tree[node];
        }

        int mid = (start + end) / 2;
        int left_child = query(node * 2, start, mid, left, right);
        int right_child = query(node * 2 + 1, mid + 1, end, left, right);

        return left_child + right_child;
    }

    public void update(int node, int start, int end, int index, int diff) {
        if (index < start || end < index) {
            return;
        }

        tree[node] += diff;

        if (start != end) {
            int mid = (start + end) / 2;
            update(node * 2, start, mid, index, diff);
            update(node * 2 + 1, mid + 1, end, index, diff);
        }
    }
}

class Solution {
    public static void main(String[] args) {
        int[] A = {1, 2, 3, 4, 5};
        SegmentTree segmentTree = new SegmentTree(A);
        int rootNode = 1;
        int start = 0;
        int end = A.length - 1;
        segmentTree.build(rootNode, start, end);
        System.out.println(segmentTree.query(rootNode, start, end, 2, 4));
        segmentTree.update(rootNode, start, end, 3, 10);
        System.out.println(segmentTree.query(rootNode, start, end, 2, 4));
    }
}

기억에 남는 심화 활용법

Linked List와 비슷한 형태로 활용 가능하다. base를 기준으로 좌측 범위에서 최댓값을 찾으면 base의 바로 왼쪽 원소이며 우측 범위에서 최솟값을 찾으면 base의 오른쪽 원소로 사용한다.
유효한 값들의 개수를 cnt로 저장하게 되면 그것이 해당 원소의 인덱스가 된다.

즉 노드에 3가지 구간정보를 저장하면 된다.
left = 살아있는 원소중 가장 왼쪽 원소의 노드 -> 왼쪽 원소 찾기
right = 살아있는 원소중 가장 오른쪽 원소의 노드 -> 오른쪽 원소 찾기
cnt = 구간에 살아있는 노드의 개수 -> 인덱스 계산

ArrayList의 경우 갱신이 O(N)
LinkedList의 경우 접근이 O(N)

그러나 해당 방법은 O(logN)으로 모든 쿼리를 해결할 수 있다.

Lazy Segement Tree

구간을 업데이트하는 쿼리가 있다면 해당 세그먼트 트리를 고려해 볼 수 있다. 기존의 세그먼트 트리 배열의 update가 배열 원소 하나의 값을 업데이트하는 것이기에 만약 범위를 모두 업데이트한다면 O(범위 * log N)이 된다. 이를 개선해보자.

아이디어

굳이 갱신할 필요가 없는 노드는 lazy라는 값을 통해 나중에 해당 노드를 방문할 일이 생길 때까지 갱신을 미룬다.

위 이미지처럼 4~10 구간에 + 10 이라는 업데이트가 발생했다면, 해당 구간에 해당하는 노드는 업데이트하고 그 자식들은 바로 업데이트하는 것이 아닌 Lazy 값을 추가해 변경이 필요함을 기록한다.
여기서 4-5노드와 6-10노드가 완전히 포함되는 구간으로 그 자식들에 Lazy 값이 추가된다.

이후 lazy 값이 붙어닜는 노드를 만나면 먼저 해당 Lazy값을 노드에 반영한다.

성능

구간 업데이트 시간복잡도 : O(logN)

구현

select연산과 비슷하게 update를 수행한다.

여기서 lazy값을 업데이트하는 과정은 다음과 같다.

• lazy값이 0이 아니라면 이를 노드에 반영한다. 이때 이 값은 (end - start + 1) = 구간의 범위에 lazy값을 곱한 값이다.
• 이후 lazy값을 자식노드에 전파한다.
• 자신의 lazy값은 0으로 변경한다.
• 만약 lazy값이 0이라면 이를 반영하지 않는다.

    static void update_lazy(long[] tree, long[] lazy, int node, int start, int end) {
    	// lazy 값이 0이 아니라면
        if (lazy[node] != 0) {
        	// 업데이트
            tree[node] += (end-start+1)*lazy[node];
            // 전파
            if (start != end) {
                lazy[node*2] += lazy[node];
                lazy[node*2+1] += lazy[node];
            }
            lazy[node] = 0;
        }
    }
	
    static void update_range(long[] tree, long[] lazy, int node, int start, int end, int left, int right, long diff) {
        update_lazy(tree, lazy, node, start, end);
        if (left > end || right < start) {
            return;
        }
        // 범위 갱신
        if (left <= start && end <= right) {
            tree[node] += (end-start+1)*diff;
            if (start != end) {
                lazy[node*2] += diff;
                lazy[node*2+1] += diff;
            }
            return;
        }
        
        update_range(tree, lazy, node*2, start, (start+end)/2, left, right, diff);
        update_range(tree, lazy, node*2+1, (start+end)/2+1, end, left, right, diff);
        tree[node] = tree[node*2] + tree[node*2+1];
    }

추천 자료

https://book.acmicpc.net/ds/segment-tree
https://book.acmicpc.net/ds/segment-tree-lazy-propagation

참고

https://yoongrammer.tistory.com/103
https://4legs-study.tistory.com/128