세그먼트 트리 - BOJ 2042

 세그먼트 트리는 주어진 array A[1 .. N] 에 대해서,

A[i..j]의 구간합을 구하는 데에 요긴하게 쓰인다.


그러면 펜윅트리랑 차이점은 무엇이냐? 펜윅트리는 구간합만 구하는데 쓰이지만 

세그먼트 트리는  update와 sum을 잘 개선하면 구간 max, 구간 min, 혹은 binary-search를 섞어 특정 k값보다 큰 원소의 갯수 등을 구할 수 있다는 장점이 있다.

세그먼트 트리를 구현하는 데에 재밌는 점은 함수를 이용하여 구현한다는 점이다.

특히 함수의 인자로 (index, start, end)가 들어가게 되는데 A[start ... end]의 구간 값을 의미하는 것이  tree[index]를 뜻하게 된다.


다음 예시 코드를 보자.


1
2
3
4
5
6
7
ll init(int index, int start, int end) {
    if(start == end) {
        return tree[index] = num[start];
    }
    int mid = (start + end/ 2;
    return tree[index] = init(2 * index, start, mid) + init(2 * index + 1, mid + 1end);
}
cs


ll은 typedef long long ll로 정의된 상태라 가정하자.

먼저 initialization의 인자로 index, start와 end가 들어가게 된다.

이후 이분 검색처럼 쪼개지면서 tree[index]에 initial value를 넣게 되는데, start와 end가 같게 되는 순간 tree[index] = arr[start]를 넣어 초기값을 만들어 준다.

당연히 tree[index / 2] 는 tree[index] + tree[index + 1]와 같은 값이 되겠다. (혹은 tree[index - 1] + tree[index])


그 후, 구간 합을 구하는 함수를 구현하자.


1
2
3
4
5
6
ll find(int index, int start, int endint left, int right){
    if (left <= start && end <= right) return tree[index];
    if (right < start || end < left) return 0;
    int mid = (start + end/ 2;
    return find(2 * index, start, mid, left, right) + find(2 * index + 1, mid+1end, left, right);
}    
cs

이번엔 추가 parameter로 left와 right가 들어간다.

left와 right가 의미하는 바는 arr[left ... right]의 구간 합을 구하겠다는 의미다.

start와 end는 현재 index가 arr[start... end]를 의미한다는 것이므로 혼동에 유의하기 바란다.


2번째 라인에서 현재 구하려는 구간 합에서 현재 index가 가르키는 합을 포함하고 있다면 당연히 현재 index가 가르키는 합을 return해주면 된다. 반면 포함하는 범위가 아니라면 0을 return 해주면 된다.


마지막으로 update를 구현하자.


1
2
3
4
5
6
7
8
9
10
11
12
void update(int index, int start, int endint t, ll diff) {
    if(start <= t && t <= end) {
        tree[index] += diff;
    }
    else {
        return;
    }
    if(start == endreturn;
    int mid = (start + end/ 2;
    update(2 * index, start, mid, t, diff);
    update(2 * index + 1, mid + 1end, t, diff);
}
cs


현재 arr[t]에 대해 값을 변경해야 할경우 변경해야 하는 변화량 diff와, 위치 t를 추가 파라미터로 설정한다.

현재 index의 start와 end가 t를 포함하면 당연히 tree[index]를 변화량만큼 더해주면 되겠다.


--------------------------------------------

재밌는 점은, initialization이 특이한 케이스의 문제가 있다.

7578) 공장이 그러한데, 공장 문제 같은 경우에는 하나하나 넣어주며 구간합을 구해야하는 케이스이다.


이 뿐만 아니라  lazy propagation과 같이 update를 느리게 하는 경우가 필요한데, 그 때는 update를 한 index만 해줄 때가 아니라 여러 index를 해줄 때 필요하다.


귀찮으므로 포스팅 끝

전혀 친절하지 않은 포스팅이 이 블로그 특징이다. ㅋ.ㅋ

댓글