구간 합을 빠르게: 펜윅 트리 (BIT)

누적합만으로는 부족할 때

구간 합 문제를 처음 만나면 다들 누적합(prefix sum)부터 떠올린다. S[i] = A[1] + ... + A[i]를 미리 구해두면 S[r] - S[l-1]로 구간 합이 바로 나온다. 여기까지는 쉽다.

근데 값이 중간에 바뀌기 시작하면 얘기가 달라진다. A[i] 하나를 갱신하는 순간 그 뒤 누적합이 전부 어긋나버리고, 다시 계산하려면 $O(N)$이 든다. 갱신이 한두 번이면 그냥 다시 계산해도 되지만, 갱신과 조회가 번갈아 수만 번씩 나오는 문제라면 이 방식으론 시간 초과를 피할 수 없다.

이럴 때 쓰는 게 세그먼트 트리나 펜윅 트리다.

핵심은 비트 연산 하나

여기서 계속 등장하는 연산이 i & -i다. 흔히 lowbit(i)라고 부르는데, $i$를 이진수로 썼을 때 가장 오른쪽에 있는 1비트만 남기고 나머지를 다 지운 값이다.

i = 6  = 0b0110
-i     = 0b1010   (2의 보수)
i & -i = 0b0010 = 2

왜 하필 i & -i로 이게 뽑히는지는 2의 보수 표현을 떠올리면 된다. -i는 ~i + 1인데, i의 최하위 1비트 아래는 전부 0이라서 1을 더해도 그 자리까지는 영향이 없고, 최하위 1비트만 그대로 살아남는다. 나머지 상위 비트는 i와 ~i가 정반대라 & 하면 다 꺼진다. 결과적으로 딱 그 1비트만 남는다.

이 lowbit(i)이 바로 "인덱스 $i$가 책임지는 구간의 길이"다. 트리 배열 tree에서 tree[i]는 A[i - lowbit(i) + 1]부터 A[i]까지의 합을 담고 있다. tree[6]은 lowbit(6)이 2니까 A[5], A[6] 두 칸의 합이고, tree[8]은 lowbit(8)이 8이라 A[1]부터 A[8]까지 한꺼번에 들고 있다. 인덱스마다 담당 구간 길이가 다르고, 그 길이가 항상 2의 거듭제곱이라는 게 이 구조의 전부라고 봐도 된다.

갱신은 위로, 조회는 아래로

말로는 잘 와닿지 않아서, 실제로 손으로 따라가 보는 게 제일 빠르다.

구현

int N;
long long tree[100010];
 
// A[i] += delta
void update(int i, long long delta){
    for(; i <= N; i += i & -i) tree[i] += delta;
}
 
// A[1] + ... + A[i]
long long query(int i){
    long long sum = 0;
    for(; i > 0; i -= i & -i) sum += tree[i];
    return sum;
}
 
// A[l] + ... + A[r]
long long rangeQuery(int l, int r){
    return query(r) - query(l - 1);
}

구현 코드는 위와 같다.

왜 겹치지 않고 딱 맞아떨어질까

여기까지 보고 나면 한 가지가 계속 걸린다. query에서 구간들을 모을 때 어떻게 겹치는 부분 하나 없이 딱 맞아떨어지는 걸까. 답은 결국 이진수에 있다.

i -= lowbit(i)는 $i$의 최하위 1비트 하나를 꺼버리는 연산이다. 13 = 0b1101로 한번 따라가 보자.

13 = 0b1101 → tree[13]: A[13..13] (lowbit=1)
12 = 0b1100 → tree[12]: A[9..12]  (lowbit=4)
 8 = 0b1000 → tree[8]:  A[1..8]   (lowbit=8)
 0 → 종료

[13,13], [9,12], [1,8]을 이어붙이면 정확히 [1,13]이 되고 겹치는 곳이 없다. 13의 이진수에서 켜져 있는 비트가 딱 세 개(1101)였고, 그 비트 하나가 꺼질 때마다 구간 하나씩을 담당한 셈이다. 결국 이 과정은 $i$의 이진수 자릿수만큼만 반복되니까 $O(\log N)$이 나온다. update가 반대 방향으로 올라가는 것도 같은 논리다.

시간복잡도

update, query 둘 다 이진수 한 자리씩 움직이니 $O(\log N)$이고, 트리 배열이 원본과 크기가 같으니 공간은 $O(N)$이다. 원소를 하나씩 update로 채우면 $O(N \log N)$이 되긴 하는데, 원본 배열이 이미 있다면 $O(N)$에 트리를 통째로 세우는 방법도 있다. 다만 코딩테스트 수준에서는 그 차이를 체감할 일이 거의 없어서, 위 구현 정도면 충분하다.

정리

• 언제: 값이 계속 바뀌는 배열에서 구간 합을 반복해서 조회할 때
• 어떻게: tree[i]에 [i - lowbit(i) + 1, i]의 합을 저장, update는 i += lowbit(i), query는 i -= lowbit(i)
• 복잡도: 갱신·조회 전부 $O(\log N)$

세그먼트 트리로도 같은 걸 풀 수 있고, 펜윅은 상수가 적어서 빠르고 가볍지만 뺄셈으로 복원 가능한 연산(합, XOR 등)에서만 쓸 수 있다. min/max처럼 역원이 없는 연산은 세그먼트 트리가 필요하다. 세그먼트 트리의 경우 모노이드 형태여야 한다.