펜윅 트리
그리고 차분 배열 트릭
- 농부 후안은 바리스타입니다 문제를 풀고 싶은데 세그먼트 트리로는 메모리 초과가 발생한다
- 세그먼트 트리보다 메모리를 덜 사용하는 펜윅 트리에 대해 간단히 정리하자
- 원소의 개수를 $N$이라 할 때 세그먼트 트리는 대략 $4N$ 크기의 배열이 필요하지만 펜윅 트리는 $N+1$ 크기의 배열이면 충분하다
- 인덱스가 $1$부터 시작해서 $N+1$ 크기의 공간이 필요하다
- 펜윅 트리는 배열의 원소를 변경하는 작업과 누적 합 계산을 $O(\log N)$에 할 수 있다
- 누적 합을 계산하는 기본 원리로 이진수를 사용한다
- 임의의 수는 $2$의 거듭제곱 꼴의 합으로 나타낼 수 있다
- 예컨대 $i=15$라면 $i=2^3+2^2+2^1+2^0$이며 이진법으로 나타내면 $1111$이다
- $i$의 최하위 비트를 $LSB(i)$라 할 때 펜윅 트리에서 $i$번째 노드 $T[i]$를 $A[i - LSB(i) + 1] + \cdots + A[i]$로 정의하자
- 즉, 배열 $A$의 $i$번째 원소부터 시작해 앞으로 $LSB(i)$개만큼 더한 것이다
- $1$번째 원소부터 $i$번째 원소까지의 누적 합을 계산하고 싶다고 해보자
- 즉, $A[1] + A[2] + \cdots + A[i]$를 계산하고 싶다
- 그런데 $i$를 $2$의 거듭제곱 꼴의 합으로 나타낼 수 있다고 했다
- $T[i]$는 정의상 누적 합의 마지막 항인 $A[i]$부터 시작해 앞으로 $LSB(i)$개만큼 더한 것이다
- 그러면 남은 누적 합은 $i' = i - LSB(i)$라 할 때 정의상 $T[i']$과 같다
- 예컨대 $i=14$라면 $i=2^3+2^2+2^1$이니까 $A[14]$에서 시작해 $2^1$개만큼 앞으로 더하고 남은 건 $2^3+2^2(=i')$이 된다
- 그럼 현재 고려하고 있는 $i$에 대해 $T[i]$를 누적 합에 더해주고 $i$를 $i-LSB(i)$로 갱신한 뒤 이를 반복하면 된다
- $i$가 $0$이 되면 반복을 멈추고 계산한 누적 합을 반환하면 끝이다
def query(tree, i):
prefix_sum = 0
while i > 0:
prefix_sum += tree[i]
i -= LSB(i)
return prefix_sum
- 만약 구간 $(A,B)$에 대해 구간 합을 계산하고 싶다면 $\operatorname{query}(\operatorname{tree}, B) - \operatorname{query}(\operatorname{tree}, A-1)$을 수행하면 된다
- $A=1$이면 누적 합이므로 단순히 $\operatorname{query}(\operatorname{tree}, B)$를 실행하면 된다
- $A[k]$에 $v$만큼 더해줬다고 해보자
- 그럼 펜윅 트리의 노드에 대해 $A[k]$를 합에 포함하면 값을 갱신해줘야 한다 (노드에 $v$만큼 더해줘야 함)
- $T[i]$가 다루는 범위는 $i - LSB(i) + 1$부터 $i$까지이다
- 즉, $i - LSB(i) + 1 \le k \le i$를 만족하는 노드의 값을 $v$만큼 증가시켜야 한다
- $i \ge k$이므로 위의 부등식을 만족하는 가장 작은 $i$는 $k$이다
- 따라서 $T[k]$에 $v$를 더해준다
- 이제 더 큰 $i$를 찾아야 하는데 어떻게 찾을 수 있을까?
- 펜윅 트리의 정의상 노드가 담당하는 구간의 크기는 $2$의 거듭제곱 꼴이다
- 즉, 펜윅 트리의 노드는 $2$의 거듭제곱 꼴 크기를 가지는 구간의 결합으로 이루어진다
- 그럼 $k$를 포함하는 마지막 구간의 크기는 $LSB(k)$이 된다
- 즉, $k=x + 2^j$로 표현될 때 $LSB(k) = 2^j$이다
- 그럼 $k$를 포함하는 마지막 구간의 크기는 $2^j$이고 이보다 더 큰 구간의 크기는 최소 $2^{j+1}$이다
- 이를 위해 $LSB(k)$를 $k$에 더해주면 준다
- 즉, $k$를 포함하는 마지막 구간의 크기를 기존 구간보다 크게 만들기 위해 최소 $LSB(k)$를 더해주는 것이다
- $k$를 $k+LSB(k)$로 갱신해가며 $T[k]$에 $v$만큼 증가시키면 된다
- 이때 $k$가 $N$을 초과하면 반복을 멈춘다
def update(tree, i, v):
while i <= N:
tree[i] += v
i += LSB(i)
- 초기에 $i=k$로 설정하여 $\operatorname{update}(\operatorname{tree}, i, v)$를 실행하면 된다
- 그런데 $LSB(i)$는 어떻게 찾을 수 있을까?
- $i$의 비트를 반전시킨 뒤 $1$을 더한 것을 $i'$이라 하자
- $i$와 $i'$에 AND 연산을 수행하면 $LSB(i)$가 된다
- $i$를 이진수로 나타내면 $xx\dots x100\dots0$과 같이 나타낼 수 있다
- $\sim i$는 $x'x'\dots x'011\dots1$이 된다 ($x'$은 $x$의 반전)
- 그럼 $\sim i + 1$은 $x'x'\dots x'100\dots0$이 되므로 $i \text{ \& } (\sim i + 1)$은 $LSB(i)$가 된다
- 이때 $-i =\; \sim i + 1$이므로 $LSB(i) = i \text{ \& } -i$이다
- 이에 대해선 $2$의 보수를 공부하자
def LSB(i):
return i & -i
- 펜윅 트리를 초기화할 때 $1$부터 $N$까지 $\operatorname{update}(\operatorname{tree}, i, A[i])$를 수행하자
- 처음에 $\operatorname{tree}$는 $0$으로 구성된 $N+1$ 크기의 배열이다
- 이러면 시간 복잡도는 $O(N\log N)$인데 $O(N)$에 할 수도 있다
- 이는 맨 위의 참고 링크의 내용을 확인하자
def init(array, n):
tree = [0] * (n + 1)
for i, a_i in enumerate(array, start=1):
update(tree, i, a_i)
return tree
- 점 업데이트, 구간 쿼리가 아닌 구간 업데이트, 점 쿼리를 수행해야 할 수 있다
- 농부 후안은 바리스타입니다 문제가 그렇다
- 이때 $D[i] = A[i] - A[i-1]$로 정의하자 (단, $D[1] = A[1]$)
- 그럼 $A[x] = D[1] + D[2] + \cdots + D[x]$가 된다 (Amazing~)
- 즉, 구간 쿼리를 수행해 $A[x]$를 계산할 수 있다
- $[L, R]$에 $v$를 더한다고 해보자
- 이는 $D[L]$에 $v$를 더하고 $D[R+1]$에 $v$를 뺀 것과 같다 (단, $R < N$)
- 구간 업데이트와 점 쿼리를 차분 배열 트릭을 사용해 점 업데이트와 구간 쿼리로 바꾼 것이다
- 여기서 차분은 시계열 자료에서 사용하는 차분과 같다