알고리즘 - 인덱스 트리(Index Tree)

개념

이진 트리 중 하나로 구간 합을 구하는데 사용된다. 말단 노드에는 원소들이 들어 있고 부모는 자식들의 합이 된다. 부모는 모든 자식들의 합이 되기 때문에 연속되는 구간의 합을 구할 수 있다. 또한 특정 위치의 값을 지속적으로 바꿔야하는 경우에도 사용된다.

https://blog.hexabrain.net/246

알고리즘 2-2강. 탐색 알고리즘 - 이진 탐색(Binary Search)

 

조건

• 부분 합을 계속해서 구해야 한다.
• 특정 인덱스 변경이 지속적으로 일어 날 수 있다.

사용 하면 좋은 경우

1. 부분 합을 계속 구해야 하는 경우
   • list에서 부분 합 배열을 구하는 경우 O(MN)의 시간이 필요하지만 Index Tree를 이용 할 경우 O(MlogN)의 시간으로 단축 시킬 수 있다.
2. 특정 인덱스의 변경이 지속적으로 일어 날 경우
   • 부분 합 배열에서 특정 인덱스를 고칠 경우 그 인덱스를 포함하는 모든 합을 다시 구해야 하기 때문에 O(MN)의 시간이 필요한데 Index Tree를 이용 할 경우 O(MlogN)으로 시간을 단축 시킬 수 있다.

구현 방법

기본 이진 탐색

1. 트리 사이즈 정하기

tydef long long ll;

ll getIndex(int num) { //num은 원소의 갯수
    ll tmp = 1;
    while(tmp < num) {
        tmp <<= 1;
    }
    return tmp * 2;
}

트리의 말단 노드에 모든 원소를 배치해야 하고 부모까지 배치해야하기 때문에 트리 배열의 사이즈는 $2^n$ ≥ (원소의 개수) 가 되는 n을 찾고 그거의 2배를 해야 부모 노드 까지 배치 할 수 있다.

1. 노드 값 채우기

ll index = getIndex(n);
ll p = index/2;
for(int i = 0; i < n; i++) { //말단 노드에 원소 입력
    scanf("%lld", &t[p]);
    p++;
}
for(ll i = (index/2) - 1; i > 0; i--) {
    t[i] = t[i*2] + t[i*2 + 1]; //왼쪽 자식과 오른쪽 자식의 합
}

트리의 말단 노드에 원소들을 입력 받고 부모 노드들에 자식들의 합을 대입한다. 부모 노드는 아래에서 부터 대입해야 한다.

b 위치의 값을 c로 바꾸기

ll idx = (index/2) + (b-1);
ll pp = idx;
if(t[idx] >= c) {
  ll gap = t[idx] - c; //바꾸려는 값과 원래의 값 차이 
  t[idx] -= gap; //차이만큼 빼기
  pp >>= 1; // 부모 노드로 가기
  while(pp > 0) { //최상단 노드(1번째 인덱스)까지 업데이트
    t[pp] -= gap;
    pp >>= 1;
  }
}
else {
  ll gap = c - t[idx];
  t[idx] += gap;
  pp >>= 1;
  while(pp > 0) {
    t[pp] += gap;
    pp >>= 1;
  }
}

특정 말단 노드의 값을 바꾸면 해당 노드의 부모 노드들의 값을 같이 바꿔준다.

구간 합 출력하기 (b~c번째 인덱스)

ll left = (index / 2) + (b - 1);
ll right = (index / 2) + (c - 1);
ll answer = 0;
while(left <= right) {
    if(left % 2 == 0) { //왼쪽에 있는 노드가 부모기준에 왼쪽에 있을때
        left /= 2; // 부모 노드로 이동 
    }
    else{//부모 기준 오른쪽에 있을때
        answer += t[left]; 
        left ++; // 해당 값을 더해주고 오른쪽 옆으로 이동
        if(left == right) { //왼쪽과 오른쪽이 만나면 해당 값 더해주고 종료
            answer += t[left];
            break;
        }
        else { //왼쪽과 오른쪽이 다르면 왼쪽의 부모로 이동
            left /= 2;
        }
    }
    if(right % 2 == 0) { //오른쪽 노드가 부모 기준 왼쪽에 있을 때
        answer += t[right]; 
        right --; //해당 값을 더해주고 왼쪽 옆으로 이동
        right /= 2; //새로운 노드의 부모로 이동
    }
    else {
        right /= 2; //부모 노드로 이동
    }
}

노드의 위치를 통해 부모 노드로 값을 구할 수 없으면 정답에 더해주고 부모 노드에 포함 된다면 값을 더하지 않고 부모 노드로 이동한다.

전체 코드

int main(int argc, const char * argv[]) {
    vector<ll> t(8194304 + 1, 0);
    scanf("%d%d%d", &n, &m, &k);
    ll index = getIndex(n);
    ll p = index/2;
    for(int i = 0; i < n; i++) {
        scanf("%lld", &t[p]);
        p++;
    }
    for(ll i = (index/2) - 1; i > 0; i--) {
        t[i] = t[i*2] + t[i*2 + 1];
    }
    for(int i = 0; i < m + k; i++) {
        scanf("%d%lld%lld", &a, &b, &c);
        if(a == 1) {
            ll idx = (index/2) + (b-1);
            ll pp = idx;
            if(t[idx] >= c) {
                ll gap = t[idx] - c;
                t[idx] -= gap;
                pp >>= 1;
                while(pp > 0) {
                    t[pp] -= gap;
                    pp >>= 1;
                }
            }
            else {
                ll gap = c - t[idx];
                t[idx] += gap;
                pp >>= 1;
                while(pp > 0) {
                    t[pp] += gap;
                    pp >>= 1;
                }
            }
        }
        else if(a == 2) {
            //구간 합 출력
            ll left = (index / 2) + (b - 1);
            ll right = (index / 2) + (c - 1);
            ll answer = 0;
            while(left <= right) {
                if(left % 2 == 0) {
                    left /= 2;
                }
                else{
                    answer += t[left];
                    left ++;
                    if(left == right) {
                        answer += t[left];
                        break;
                    }
                    else {
                        left /= 2;
                    }
                }
                if(right % 2 == 0) {
                    answer += t[right];
                    right --;
                    right /= 2;
                }
                else {
                    right /= 2;
                }
            }
            printf("%lld\\n", answer);
        }
    }
    return 0;
}

시간 복잡도

• O(MlogN)