[백준 2042번] 구간 합 구하기 - Python

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

2042번: 구간 합 구하기

풀이

사용한 알고리즘 : BIT(Binary Index Tree)

 

풀이전략

BIT 알고리즘에 대해 이해하고 있어야만 이 문제를 풀 수 있다.

따라서 BIT 알고리즘에 대해 설명하도록 하겠다.

 

1. 트리 구조 만들기

원소가 8개인 리스트가 주어졌을 때, 해당 리스트를 바이너리 인덱스 트리로 만들어 보겠다.

1~8 사이의 숫자들을 이진수로 바꾸면

1 = 0001(2)    2 = 0010(2)    3 = 0011(2)    4 = 0100(2)

5 =  0101(2)    6 = 0110(2)      7 = 0111(2)     8 = 1000(2) 이다.

 

여기서 각 수를 이진수로 만든 비트를 자세히 보고, 비트의 가장 끝에 있는 1의 위치를 주목해보자.

1, 3, 5, 6과 같은 경우 마지막 1이 0001(2) 자리에 위치하고 있다. 

이는 1로서 해당 자리는 그 자리에 있는 1개의 합만을 나타낸다.

 

근데 2, 6은 마지막 1이 0010(2) 자리에 위치하고 있다.

이는 2로서 해당 자리는 그 자리에 있는 수 포함 2개의 합을 나타낸다.

 

4는 마지막이 1이 0100(2) 자리에 위치하고 있다.

이는 4로서 해당 자리는 그 자리에 있는 수 포함 4개의 합을 나타낸다.

 

마지막으로 8은 마지막 1이 1000(2) 자리에 위치하고 있다.

이는 8로서 해당 자리는 그 자리에 있는 수 포함 8개의 합을 나타낸다.

 

2. 값을 변경할 때

만약 3번 인덱스에 있는 값을 변경한다면, 위와 같이 3, 1~4합, 1~8합을 수정해야 한다.

그런데 어떻게하면 이 트리를 활용하여 빠르게 변경해야하는 값들의 인덱스를 찾을 수 있을까?

이는 이진수에서 보수를 활용하여 구할 수 있다.

 

3 = 0011(2) 로 2의 보수를 구하면 1100(2) + 0001(2) = 1101(2) 이다. (역수에 1을 더한 값이므로)여기서 3 & -3 = 0011(2) & 1101(2) = 0001(2)가 나오는데, 특정 수와 특정 수의 보수를 and 연산하면 자신의 비트 중 가장 끝에 있는 1의 위치를 찾을 수 있게 된다.

 

이를 활용하여 위 문제에 활용해보면k = k + (k & -k) 로 k번째 인덱스의 값이 변했을 경우 바꿔야 할 인덱스들을 구할 수 있다.위의 문제를 예로 들면 

1) k =3 일 때,  0011(2) + ( 0011(2) & 1101(2) ) = 0011(2) + 0001(2) = 0100(2) = 4 

2) k =4 일 때,  0100(2) + ( 0100(2) & 1100(2) ) = 0100(2) + 0100(2) = 1000(2) = 8

로 우리가 구하려고 했었던 3, 4, 8이 구해지는 것을 알 수 있다.

 

따라서 값을 변경할 때, 위 방법을 통해 인덱스를 바꿔가면서 원래 수는 빼고 변하는 수를 더해주면 업데이트 시켜줄 수 있다.

 

3. 누적합 구하기

만약 인덱스 7의 값을 변경하려고 한다면,

위의 그림처럼 1~4 번째 값의 합을 저장하는 인덱스 4, 5~6 번째 값의 합을 저장하는 인덱스 6, 7번 인덱스를 구할 수 있어야 한다.

 

이 역시도 2의 보수를 활용하면 쉽게 구할 수 있다.

이번엔 k = k - ( k & -k)를 구하면 된다.

 

1) k=7일 때, 0111(2) - ( 0111(2) & 1001(2) ) = 0111(2) - 0001(2) = 0110(2) = 6

2) k=6일 때, 0110(2) - ( 0110(2) & 1010(2) ) = 0110(2) - 0010(2) = 0100(2) = 4

로 우리가 구하려고 했었던 7, 6, 4가 구해지는 것을 알 수 있다.

 

따라서 해당 인덱스들에 대응되는 값들을 모두 더하면 누적합을 구할 수 있다.

 

아래 제출 코드는 이를 활용하여 해결하였다.

 

제출코드

import sys

input=sys.stdin.readline

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

def update(x,val):
    while x<=n:
        tree[x]+=val
        x+=(x&(-x))

def interval_sum(start,end):
    return prefix_sum(end)-prefix_sum(start-1)

n,m,k=map(int,input().split())

arr=[0]*(n+1)
tree=[0]*(n+1)

for i in range(1,n+1):
    arr[i]=int(input())
    update(i,arr[i])

for i in range(m+k):
    a,b,c=map(int,input().split())
    if a==1:
        update(b,c-arr[b])
        arr[b]=c
    else:
        print(interval_sum(b,c))