CS/자료구조

세그먼트 트리

Meluu_ 2025. 7. 9. 10:12

배열의 연속된 구간의 합을 O(Log N)시간만에 할 수 있는 자료구조

합 뿐만 아니라 구간 최대 최소, 곱 등도 구할 수 있다.

 

세그먼트 트리는 보통 3가지로 구성된다

init : 세그먼트 트리 초기화

query : 원하는 구간의 값을 구함

update : 배열에서 요소 a를 변경했을 때 값을 갱신함

 

 

트리 크기

세그먼트 트리도 결국 이진 트리 기반이기에 가질 수 있는 노드는

2^h 개가 된다.

 

h = ceil(log2(n)) , n = 배열의 요소 개수 

treeSize = 2^h 

 

높이를 올리는 이유

n = 9 일때 log2(9) = 3.xx 로 나온다.

그렇다고 높이를 3으로 해버리면

2^3 이 되어 트리 크기가 8이 되고

 

9를 넣을 공간이 없어진다. 

그래서 높이를 올림(ceil)해서 4로 잡아야 안전하다.

 

 

코드에는 1 << (H+1) 로 표현했다. 

(2 << H ) == (1 << H + 1 )

 

 

해당 자료구조를 사용하려할때 생각해봐야할 것

  • 누적합 필요 여부
  • 구간 최대/최솟값 구하기
  • 갱신이 자주 일어나는지
  • 구간 관련된 여러 문제
  • Lazy Propagation (계산을 나중으로 미룸)

 

기본 적인 누적합 세그먼트 트리를 구현해보기

public class 세그먼트트리구현해보기 {

    static int n, size;
    static int[] tree, arr;

    public static void main(String[] args) {

        n = 10;
        arr = new int[] {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

        int H = (int) Math.ceil(Math.log(n) / Math.log(2));
        size = 1 << (H + 1);
        tree = new int[size];

        init(1, 1, n);

        int query = query(1, 1, n, 3, 5);

        System.out.println(query);
        System.out.println(query(1, 1, n, 1, 10));

        int diff = 2 - arr[5];
        arr[5] = 2;
        
        update(1, 1, n, 5, 2 - 5);
        System.out.println(query(1, 1, n, 3, 6));

    }

    /**
     * 처음 값 초기화
     * @param treeNode 트리노드 값
     * @param s         시작 위치
     * @param e         끝 위치 
     * @return
     */
    private static int init(int treeNode, int s, int e) {
        // 리프 노드이면 값 삽입
        if (s == e) {
            return tree[treeNode] = arr[s];
        }

        int mid = (s + e) >>> 1;

        // 리프노드를 제외한 나머지는 누적합으로 구한다. 
        return tree[treeNode] = init(treeNode * 2, s, mid)
                + init(treeNode * 2 + 1, mid + 1, e);
    }


    /**
     * 세그먼트 트리에 변경된 값을 업데이트 
     * @param treeNode 트리 노드 값
     * @param s       시작 위치
     * @param e       끝 위치
     * @param idx     변경한 요소의 인덱스
     * @param diff    기존 값과 변경한 값의 차이  
     */
    private static void update(int treeNode, int s, int e, int idx, int diff) {
        // 변경한 요소의 인덱스 범위를 벗어났으면 업데이트 할 필요 X
        if (idx < s || e < idx) return;
        
        // 변경 값 반영
        tree[treeNode] += diff;

        int mid = (s + e) >>> 1;

        if (s != e) {
            update(treeNode * 2, s, mid, idx, diff);
            update(treeNode * 2 +1, mid + 1, e ,idx, diff);
        }
    }

    /**
     *  쿼리로 원하는 범위의 누적합을 반환
     * @param treeNode  트리 노드 값
     * @param s         시작 위치
     * @param e         끝 위치
     * @param left      찾는 범위 시작 위치
     * @param right     찾는 범위 끝 위치 
     * @return          범위의 누적합
     */
    private static int query(int treeNode, int s, int e, int left, int right) {
        if (right < s || e < left) return 0;
        
        // 범위 내라면 해당 트리의 노드값을 반환
        if (left <= s && e <= right) return tree[treeNode];

        int mid = (s + e) >>> 1;

        return query(treeNode * 2, s, mid, left, right)
                + query(treeNode * 2 + 1, mid + 1, e, left, right);
    }
}

'CS > 자료구조' 카테고리의 다른 글

연결리스트  (0) 2024.07.09
큐(queue)와 덱(deque)  (0) 2024.07.02
스택 (Stack)  (0) 2024.07.01
배열, 구조체, 포인터  (0) 2024.06.28
순환  (0) 2024.01.26