Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

이진 인덱스 트리(BIT): C++로 구현하는 범위 업데이트와 범위 합 쿼리

문제 개요

크기가 n인 배열이 주어지며, 처음에는 모든 원소가 0으로 초기화되어 있습니다. 이 배열에 대해 다음 두 가지 종류의 쿼리를 수행해야 합니다.

  • update(l, r, value) — 인덱스 l부터 r 사이의 모든 배열 원소에 value를 더합니다. 예를 들어 update(2, 4, 5)는 인덱스 2, 3, 4의 원소에 각각 5를 더합니다.

  • getRangeSum(l, r) — 인덱스 l부터 r 사이에 있는 원소들의 합을 구합니다. 예를 들어 getRangeSum(4, 7)은 인덱스 4, 5, 6, 7에 해당하는 원소들의 합을 반환합니다.

구체적인 예제를 통해 문제를 살펴보겠습니다.

입력

n = 7 , arr[7] = {0,0,0,0,0,0,0}
Q1 = update(3, 6, 4)
Q2 = update(0, 4, 2)
Q3 = sum(2, 5)

출력

10

설명

Q1 - update(3, 6, 4) 실행 후 → {0, 0, 0, 4, 4, 4, 4}
Q2 - update(0, 4, 2) 실행 후 → {2, 2, 2, 2, 2, 4, 4}
Q3 - sum(2, 5) = 2 + 2 + 2 + 4 = 10

단순한 접근법과 그 한계

가장 직관적인 방법은 업데이트 쿼리가 들어올 때마다 해당 범위의 배열을 직접 수정하고, 합 쿼리가 들어오면 범위를 순회하며 값을 더하는 것입니다. 하지만 이 방법은 각 연산이 최악의 경우 O(n)의 시간이 걸리므로, 쿼리 개수가 많아지면 비효율적입니다. 따라서 더 효율적인 접근 방식을 살펴보겠습니다.

효율적인 접근: 이진 인덱스 트리(BIT)

이 문제는 이진 인덱스 트리(Binary Indexed Tree), 즉 펜윅 트리(Fenwick Tree)를 활용하면 각 연산을 O(log n) 시간에 처리할 수 있습니다.

먼저 범위 합 쿼리 sum[l, r]을 다음과 같이 접두사 합(prefix sum) 형태로 분해합니다.

sum[l, r] = sum[0, r] − sum[0, l−1]

즉, sum[0, k] 형태의 쿼리만 효율적으로 처리할 수 있다면 범위 합도 쉽게 구할 수 있습니다. 이제 업데이트 쿼리 update(l, r, value)가 sum[0, k]에 미치는 영향을, k가 업데이트 범위 [l, r]과의 상대적 위치에 따라 세 가지 영역으로 나누어 분석해 보겠습니다.

  • 영역 1 — k가 l보다 앞선 경우 (k < l)
    업데이트 쿼리가 sum[0, k]에 아무런 영향을 주지 않습니다.

  • 영역 2 — k가 범위 내부에 있는 경우 (l ≤ k ≤ r)
    sum[0, k]는 l부터 k까지 원소들의 증가분을 반영하며, 정확히 value × (k − l + 1)만큼 증가합니다.

  • 영역 3 — k가 r보다 큰 경우 (k > r)
    sum[0, k]는 l부터 r까지 전체 범위의 증가분, 즉 value × (r − l + 1)을 반영합니다.

이 세 가지 경우를 하나의 공식으로 통합하기 위해 두 개의 BIT를 사용합니다. 첫 번째 트리(BITree1)에는 차분 배열 역할을 하는 값을 저장하고, 두 번째 트리(BITree2)에는 보정항을 저장합니다. 그러면 임의의 위치 x까지의 합은 다음과 같이 계산됩니다.

sum(x) = getSum(BITree1, x) × x − getSum(BITree2, x)

이제 범위 업데이트와 범위 합 쿼리를 해결하는 프로그램을 살펴보겠습니다.

C++ 구현 예제

#include <bits/stdc++.h>
using namespace std;

// 인덱스 i까지의 접두사 합을 구하는 함수
int getSum(int BITree[], int i){
    int sum = 0;
    i++;
    while (i > 0) {
        sum += BITree[i];
        i -= i & (-i);
    }
    return sum;
}

// BIT의 인덱스 i에 val을 더하는 함수
void updateBITree(int BITree[], int n, int i, int val) {
    i = i + 1;
    while (i <= n) {
        BITree[i] += val;
        i += i & (-i);
    }
}

// 범위 [l, r]에 value를 더하는 함수
void update(int BITTree1[], int BITTree2[], int n, int l, int r, int value) {
    updateBITree(BITTree1, n, l, value);
    updateBITree(BITTree1, n, r + 1, -value);
    updateBITree(BITTree2, n, l, value * (l - 1));
    updateBITree(BITTree2, n, r + 1, -value * r);
}

// 위치 x까지의 합 계산
int sum(int x, int BITTree1[], int BITTree2[]) {
    return (getSum(BITTree1, x) * x) - getSum(BITTree2, x);
}

// 범위 [l, r]의 합 계산
int getRangeSum(int l, int r, int BITTree1[], int BITTree2[]) {
    return sum(r, BITTree1, BITTree2) - sum(l - 1, BITTree1, BITTree2);
}

// 크기 n짜리 BIT 생성 및 초기화
int *createBITree(int n) {
    int *BITree = new int[n + 1];
    for (int i = 1; i <= n; i++)
        BITree[i] = 0;
    return BITree;
}

int main(){
    int n = 7;
    int *BITTree1, *BITTree2;
    BITTree1 = createBITree(n);
    BITTree2 = createBITree(n);
    update(BITTree1, BITTree2, n, 3, 6, 9);
    update(BITTree1, BITTree2, n, 0, 4, 5);
    cout << "The output of sum query after applying all update queries is \t"
         << getRangeSum(1, 5, BITTree1, BITTree2);
    return 0;
}

출력

The output of sum query after applying all update queries is 47

이 프로그램은 update(3, 6, 9)와 update(0, 4, 5)를 차례로 적용한 뒤 getRangeSum(1, 5)을 호출합니다. 두 업데이트가 끝난 시점의 배열은 {5, 5, 5, 14, 14, 9, 9}이므로, 인덱스 1부터 5까지의 합은 5 + 5 + 14 + 14 + 9 = 47이 됩니다.

복잡도 분석

  • 시간 복잡도 — 범위 업데이트와 범위 합 쿼리 모두 각각 두 번의 BIT 연산만 수행하므로 O(log n)입니다.

  • 공간 복잡도 — 두 개의 BIT를 사용하므로 O(n)입니다.

마무리

이진 인덱스 트리는 단순한 포인트 업데이트·접두사 합 조회를 넘어, 두 개의 트리를 조합하면 범위 업데이트와 범위 합 쿼리까지 효율적으로 처리할 수 있는 강력한 자료구조입니다. 세그먼트 트리에 비해 구현이 간결하고 메모리 사용량이 적다는 장점이 있어, 코딩 테스트나 경쟁 프로그래밍에서 빈번하게 활용됩니다.