세그먼트 트리

세그먼트 트리는 point update와 range query를 모두 O(logN)의 시간에 처리할 수 있는 자료구조입니다.

세그먼트 트리는 아래처럼 생긴 이진 트리입니다. 구간의 길이가 1인(원소의 개수가 하나인) 구간은 리프 노드가 되고, 나머지 노드는 자식 노드를 합친 구간 합을 들고 있습니다.

배열의 길이가 N이라면 세그먼트 트리의 노드 개수는 2N - 1개가 됩니다.

힙과 마찬가지로 세그먼트 트리는 배열로 구현할 수 있습니다. 0번 노드는 버리고, 1번 노드가 루트가 되는 방식으로 구현하면 완전 이진 트리의 모습을 하며, 필요한 배열의 길이는 2N이 되고, 리프 노드의 인덱스는 [ N , 2N )입니다.

완전 이진 트리이기 때문에 높이는 O(logN)이 됩니다.

구현

#include<iostream>
#include<algorithm>
#define MAX_K 1000000

#define parent(i) (i >> 1)
#define left(i) (i << 1)
#define right(i) (i << 1 | 1)

using namespace std;

int arr[MAX_K];
int tree[MAX_K];

void buildTree(int i, int start, int end){
    // leaf 노드가 들어갈 위치
    if(start == end){
        tree[i] = arr[start];
        return;
    }

    int mid = (start + end) / 2;
    buildTree(left(i),start,mid);
    buildTree(right(i),mid+1,end);

    // sum tree
    tree[i] = tree[left(i)] + tree[right(i)];
    // max tree
    tree[i] = max(tree[left(i)], tree[right(i)]);
    // min tree
    tree[i] = min(tree[left(i)], tree[right(i)]);
}

int query(int i, int start, int end, int left, int right){
    // 유효한 범위가 아니므로 종료(무의미한 값 반환)
    if(start > right || end < left){
        // sum tree
        return 0;
        // max tree
        return 0;
        // min tree
        return MAX_K;
    }

    // i가 left ~ right를 모두 포함하는 범위이므로 i값 반환 
    if(left <= start && end <= right)
        return tree[i];

    int mid = (start+end) / 2;

    int q1 = query(left(i), start, mid, left, right);
    int q2 = query(right(i),mid+1, end, left, right);

    // sum tree
    int result = q1+q2;
    // max tree
    int result = max(q1,q2);
    // min tree
    int result =  min(q1,q2);
    
    return result;
}

void updateSumTree(int i, int index ,int start, int end, int diff){
    // 유효한 범위가 아니므로 종료
    if(index < start || index > end)
        return;

    // sum tree
    tree[i] += diff;

    if(start != end){
        int mid = (start + end) / 2;
        updateSumTree(left(i), index, start, mid, diff);
        updateSumTree(right(i), index, mid+1, end, diff);
    }
}

int updateMaxMinTree(int i, int index ,int start, int end, int value){
    // 유효한 범위가 아니므로 종료
    if(index < start || index > end)
        return tree[i];

    if(start == end)
        return tree[i] = value;

    int mid = (start + end) / 2;
    int left = updateMaxMinTree(left(i), index, start, mid, value);
    int right = updateMaxMinTree(right(i), index, mid+1, end, value);

    // max tree
    return tree[i] = max(left,right);
    // min tree
    return tree[i] = min(left,right);
}


int main(void){
    int len;
    cin >> len;

    for(int i=0;i<len;i++)
        cin >> arr[i];

    memset(tree,0,sizeof(int)*len*4);
    buildTree(1,0,len-1);

    while(1){
        int cmd;
        cin >> cmd;

        if(cmd == -1)
            break;

        if(cmd == 1){
            int left, right;
            cin >> left >> right;

            cout << query(1,0,len-1,left-1,right-1) << endl;
        }else if(cmd == 2){
            int i, value;
            cin >> i >> value;

            int diff;

            // sum tree
            diff = value - arr[i];
            arr[i] = value;

            updateSumTree(1,i,0,len-1,diff);

            // max & min tree            
            arr[i] = value;

            updateMaxMinTree(1,i,0,len-1,value);
        }else if(cmd == 3){
            for(int i=1;i<4*len;i++)
                cout << tree[i] << " ";
            cout << endl;
        }
    }
    return 0;
}

References

https://reakwon.tistory.com/44
https://eun-jeong.tistory.com/18