문제 54
구간 합과 쿼리
길이가 $N$인 정수 수열 $A_1, A_2, \cdots, A_N$이 주어진다.
$Q$개의 쿼리를 순서대로 처리하라.
쿼리는 다음 중 하나이다.
1 i x: $A_i$의 값을 $x$로 바꾼다. $(1 \le i \le N; 0 \le x \le 10^9)$2 l r: $A_l + A_{l+1} + \cdots + A_r$을 출력한다. $(1 \le l \le r \le N)$
입력
첫 번째 줄에 두 정수 $N$, $Q$가 공백으로 구분되어 주어진다. $(1 \le N, Q \le 200\,000)$
두 번째 줄에 $N$개의 정수 $A_1, A_2, \cdots, A_N$이 공백으로 구분되어 주어진다. $(0 \le A_i \le 10^9)$
세 번째 줄부터 $Q$개의 줄에 걸쳐 쿼리가 하나씩 주어진다.
출력
2 l r 쿼리마다 구간 합을 한 줄에 하나씩 출력한다.
예제 입력 1
5 7 1 2 3 4 5 2 1 5 1 3 10 2 2 4 1 5 0 2 4 5 2 3 3 2 1 1
예제 출력 1
15 16 4 10 1
공식 해설
구간 합을 접두사 합의 차로 구하기
첫 $r$개 원소의 합을 $P(r)$라 하면 구간 $[l,r]$의 합은 $P(r)-P(l-1)$이다. 원소가 바뀌어도 이 합을 빠르게 구하도록 Fenwick 트리를 사용한다.
양의 정수 $i$를 나누는 가장 큰 $2$의 거듭제곱을 $b(i)$라 하자. 트리의 $i$번째 칸에는 구간 $[i-b(i)+1,i]$의 합을 저장한다.
접두사 합을 구할 때는 현재 위치의 저장값을 더한 뒤 $i\leftarrow i-b(i)$로 이동한다. $i=0$이 될 때까지 반복하면 접두사를 겹치지 않는 구간들로 나누어 더한 것이다. 원소에 변화량을 더할 때는 $i\leftarrow i+b(i)$로 이동하며 $N$ 이하인 모든 위치에 변화량을 더한다. 이 위치들이 해당 원소를 포함하는 저장 구간들이다.
대입을 변화량으로 바꾸기
$A_i$를 $x$로 바꾸려면 트리에 $x$ 자체가 아니라 $\Delta=x-A_i$를 더해야 한다. 현재 원소 값도 함께 저장하여 변화량을 구한 뒤 $A_i=x$로 갱신한다.
초기에는 각 칸에 $A_i$를 놓고 작은 인덱스부터 순회한다. $i$번째 합을 $i+b(i)\le N$인 부모 위치에 더하면 완성된 구간 합이 전달되므로 선형 시간에 트리를 만들 수 있다.
각 연산은 $O(\log(N+1))$개의 위치만 방문한다. 전체 시간복잡도는 $O(N+Q\log(N+1))$, 공간복잡도는 $O(N)$이다. 구간 합은 최대 $N\cdot10^9$까지 커질 수 있다.
구현 · 언어별 풀이
C++ 구현
트리와 원래 값을 vector<long long>으로 저장한다. $b(i)$는 양의 인덱스에서 i & -i로 구할 수 있다. 갱신은 i += i & -i, 접두사 합은 i -= i & -i로 이동한다. 인덱스 $0$은 갱신 시작점으로 사용하지 않는다.
C 구현
현재 값 배열과 Fenwick tree를 long long으로 둔다. 대입 x는 delta=x-old로 바꿔 tree에 더하고 현재 값도 x로 교체한다. 초깃값은 각 노드에서 부모로 합을 전달하여 선형 시간에 만들 수 있다.
$1$부터 센 위치 i에 x를 대입하는 부분이다. 기존 값과의 차이만 더해야 전체 합이 유지된다.
long long delta = x - current[i];
current[i] = x;
for (int j = i; j <= n; j += j & -j) tree[j] += delta;
Python 구현
Fenwick tree와 현재 값을 array('q')에 저장할 수 있다. 대입을 단순히 x만큼 더하는 연산으로 혼동하지 말고 차이만 반영한다. 누적합은 prefix(r)-prefix(l-1)이다.
$1$부터 센 위치 i에 x를 대입하는 부분이다. 기존 값과의 차이만 더해야 전체 합이 유지된다.
delta = x - current[i]
current[i] = x
j = i
while j <= n:
tree[j] += delta
j += j & -j
Java 구현
long[] current와 tree를 사용한다. 대입은 delta=x-current[i]를 더하고 current를 갱신한다. int 인덱스에서 i & -i로 이동하며 합을 int로 변환하지 않는다.
$1$부터 센 위치 i에 x를 대입하는 부분이다. 기존 값과의 차이만 더해야 전체 합이 유지된다.
long delta = x - current[i];
current[i] = x;
for (int j = i; j <= n; j += j & -j) tree[j] += delta;
Rust 구현
값과 tree는 i64, 인덱스는 usize다. 양수 인덱스의 lowbit는 i & i.wrapping_neg()로 구할 수 있다. 대입은 이전 값과의 차이를 반영한다.
$1$부터 센 위치 i에 x를 대입하는 부분이다. 기존 값과의 차이만 더해야 전체 합이 유지된다.
let delta = x - current[i];
current[i] = x;
let mut j = i;
while j <= n {
tree[j] += delta;
j += j & j.wrapping_neg();
}
JavaScript 구현
원소 값과 Fenwick 트리 합은 Number로 계산한다. 전체 합의 상한 $2\cdot10^{14}$은 안전 정수 범위 안이지만 $32$비트를 넘으므로 두 배열 모두 Float64Array를 사용한다.
현재 값을 보관하고 대입 명령에서 delta = x - a[i]를 계산해 트리에 더한다. 대입을 단순 덧셈으로 처리하지 않는다. 다음은 $1$부터 시작하는 트리의 기본 연산이다.
function add(i, x) {
for (; i <= n; i += i & -i) tree[i] += x;
}
function sum(i) {
let s = 0;
for (; i > 0; i -= i & -i) s += tree[i];
return s;
}
비트 연산은 작은 인덱스에만 적용하며 합에 | 0을 적용하면 값이 잘린다. 닫힌 구간 $[l,r]$의 합은 sum(r) - sum(l - 1)이다.