SOJ ONLINE JUDGE

문제 54

구간 합과 쿼리

내 상태
미제출
난이도
54번 문제 난이도 보기
Gold I
출제자
rlatjwls7882
시간 제한
1000 ms
메모리 제한
512 MB
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)이다.

제출

편집기에서 나가려면 Esc를 누른 뒤 Tab 또는 Shift와 Tab을 누르세요.