[백준 2042][파이썬] 구간 합 구하기(펜윅 트리)

https://www.acmicpc.net/problem/2042

2042번: 구간 합 구하기

https://growth-coder.tistory.com/162

이전 포스팅에서 펜윅 트리(BIT, 이진 인덱스 트리)의 사용법에 대해서 공부하였다.

 

이 펜윅 트리를 사용하면 구간 합 문제를 쉽게 구현할 수 있다.

 

우선 갱신 함수이다.

def update(i, num):
    arr[i]+=num
    while i<=n:
        tree[i]+=num
        i+=i&-i

이 갱신 함수는 갱신에만 사용하는 것이 아니라 펜윅 트리의 생성에도 사용할 함수이다.

 

우선 펜윅 트리 생성부터 보자.

 

가장 먼저 arr라는 배열에는 모든 값을 0으로 저장한다.

 

그리고 입력이 들어오면 순서대로 입력한 값을 더해준다.

 

초기값이 0이므로 입력한 값 그대로 저장된다. 이렇게 구현한 이유는 갱신에도 이를 활용하기 위함이다.

 

위 배열을 보면 처음에 인덱스 1에 해당하는 값(1)을 입력받으면 인덱스 1을 포함하는 구간에 해당 수를 더해준다.

 

해당 구간은 초록색 블록을 기준으로 1, 2, 4, 8, 16번 블록에 해당한다.

 

인덱스 5에 해당하는 값(1)를 입력받으면 5, 6, 8, 16번 초록색 블록에 값을 더하는 것이다.

 

이렇게 인덱스 1부터 16까지 진행하게되면 다음과 같은 펜윅 트리가 생성된다.

그 다음 트리의 생성이 아닌 배열의 값을 바꾸는 경우를 보자.

 

이제 위에서 A[3]을 3으로 바꿨다고 해보자.

 

기존 A[3]과 -2만큼 차이가 난다. update(3, -2)을 실행하면 A[3]의 값에 2를 빼고 인덱스 3을 포함하는 구간에 모두 2를 빼준다.

이렇게 update 함수 하나로 펜윅 트리 생성과 갱신 연산을 수행할 수 있다.

 

그 다음 처음부터 특정 인덱스까지의 합을 구하는 sum 함수이다.

def sum(i):
    res=0
    while i>0:
        res+=tree[i]
        i-=i&-i
    return res

A[11]까지의 합을 구한다면 아래와 같은 블록들을 더한다는 의미이다.

만약 A[3]부터 A[11]까지의 구간 합을 구하고 싶다면 인덱스 11까지의 합에서 인덱스 2(3-1)까지의 합을 빼면 된다.

 

import sys
input = sys.stdin.readline
n, m, k = map(int, input().split())
arr=[0]*(n+1)
tree=[0]*(n+1)
def update(i, num):
    arr[i]+=num
    while i<=n:
        tree[i]+=num
        i+=i&-i
def sum(i):
    res=0
    while i>0:
        res+=tree[i]
        i-=i&-i
    return res
for i in range(1, n+1):
    update(i, int(input()))
for _ in range(m+k):
    a, b, c = map(int, input().split())
    if a==1:
        update(b, c-arr[b])
    else:
        print(sum(c)-sum(b-1))