백준 2042번: 구간 합 구하기 (Segment Tree)

문제링크:2042번: 구간 합 구하기 (acmicpc.net)

2042번: 구간 합 구하기

이 문제는 세그먼트 트리를 이용하여 구간 합을 구하는 문제이다. 구간합을 구하는 방법에는 여러가지가 있는데 만약 수의 개수가 N개인 리스트에서 구간 합을 구한다고 가정하자. 먼저 실제로 그냥 더하는 경우는 구간 합을 구하는 경우 최대 O(N)의 시간 복잡도를 가진다. 그리고 누적 합을 이용하는 경우가 있는데 이는 누적 합이 저장되어 있기에 구간 합을 구하는데 시간 복잡도가 O(1)밖에 되지 않지만 위 문제처럼 중간에 수를 바꾼다면, 누적 합 리스트를 다시 만들어야 하기에 수정하는 경우 최대 O(N)의 시간 복잡도를 가진다.

 

세그먼트 트리는 각 부분의 합을 트리 형태로 저장하여 구간 합을 구하는데 O(logN)의 시간 복잡도를 가지고, 중간에 수를 바꾸는 경우 트리를 재설정 하는데에도 O(logN)의 시간복잡도를 가진다. 그러므로 수를 바꾸고, 구간합을 구하는데 O(logN)의 시간 복잡도로 지금까지 배운 다른 방식보다 우위를 갖는다.

 

세그먼트 트리의 루트의 인덱스는 1이고, 1부터 N까지 전체의 합을 저장하고 있다. 루트의 자식은 이진 트리 방식으로 두개씩 늘어나는데, 각각 부모 노드의 합 범위의 절반씩 저장하는 방식으로 트리가 진행된다. 예시로 8개의 수로 이루어진 수열의 합을 세그먼트 트리로 표현하면,

 

인덱스 1 는 1~8까지의 합,

인덱스 2 는 1~4까지의 합, 인덱스 3은 5~8까지의 합,

인덱스 4는 1~2, 인덱스 5는 3~4, 인덱스 6은 5~6, 인덱스 7은 7~8,

인덱스 8은 1, 인덱스 9는 2, ....

 

이렇게 트리가 진행된다. 세그먼트 트리의 특징으로는 이진 트리이기 때문에 어떤 노드의 인덱스가 k라면 자식 노드의 인덱스는 항상 2k, 2k+1이라는 점이다. 이 특징을 이용하면 이진 트리를 쉽게 구현할 수 있다.

import sys
N,M,K=map(int,sys.stdin.readline().rstrip().split())

numlst=[0]
for _ in range(N):
    numlst.append(int(sys.stdin.readline().rstrip()))
tree=[0 for _ in range(4*N)]

def fill_tree(start,end,index):
    if start==end:
        tree[index]=numlst[start]
        return tree[index]
    
    mid=(start+end)//2
    tree[index]=fill_tree(start,mid,2*index) + fill_tree(mid+1,end,2*index+1)
    return tree[index]


def return_sum(start,end,left,right,index):
    if end<left or start>right:
        return 0
    
    if start<=left and right<=end:
        return tree[index]
    
    mid=(left+right)//2
    return return_sum(start,end,left,mid,2*index) + return_sum(start,end,mid+1,right,2*index+1)


def change_tree(start,end,index,changeidx):
    if start==end:
        tree[index]=numlst[start]
        return tree[index]
    
    mid=(start+end)//2
    if changeidx<= mid:
        tree[index]=change_tree(start,mid,2*index,changeidx) + tree[2*index+1]
    if changeidx > mid:
        tree[index]=tree[2*index] + change_tree(mid+1,end,2*index+1,changeidx)
    return tree[index]


fill_tree(1,N,1)

for _ in range(M+K):
    a,b,c=map(int,sys.stdin.readline().rstrip().split())
    if a==1:
        numlst[b]=c
        change_tree(1,N,1,b)
    if a==2:
        print(return_sum(b,c,1,N,1))

먼저 세그먼트 트리를 저장할 리스트 tree를 만들어 주는데, tree의 최대 인덱스는 가장 깊은 자식 노드들이 N개 이상인 완전 이진 트리의 개수만큼 만들어 주어야 한다. 완전 이진 트리의 가장 하위 노드들의 개수가 N개의 수를 담기에 충분해야 한다. 완전 이진 트리의 가장 하위 노드는 2,4,8,16,32 순으로 2^k개만큼 증가한다 그리고 부모 노드들이 (2^k)-1 만큼 또 존재하므로 2*(2^k) 만큼 길이의 리스트를 만들어 주어야 한다. N보다 큰 가장 작은 2^k의 수에 2배를 하면 딱 맞게 떨어지지만 메모리 제한이 없다면 4*N으로 잡아주면 충분하다.

 

fill_tree()함수는 세그먼트 트리를 채워주는 함수이다. start,end,index를 매개변수로 start와 end가 같아질 때까지 재귀로 구간 합의 범위를 절반으로 나눠주어 각 인덱스에 해당하는 (start,end) 범위의 구간 합을 tree[index]에 저장 해 준다.

 

다음으로 change_tree()함수는 만약 수열 중 하나를 바꿔야 할 때 사용하는 함수이다. fill_tree()와 비슷하게 절반을 나누는 과정을 재귀를 사용하는데, 만약 바꾸는 수의 인덱스인 changeidx가 해당 노드의 범위 안에 있지 않다면 값을 바꿔주지 않아도 되기에 changeidx가 범위에 있는 자식 노드만 change_tree() 함수로 재귀해 변경 해나간다.

 

마지막으로 return_sum()함수는 세그먼트 트리를 이용해 구간 합을 얻어내는 함수이다. 구해야 할 구간 합이 start, end에 들어가고 현재 index의 노드의 범위가 left, right에 저장된다.

1. 만약 (start, end) 범위가 (left,right)범위에 겹치지 않는다면 해당 index 노드의 합 범위가 구해야 할 범위 밖이므로 필요 없으므로 return 0을 해준다.

2. 만약 (start,end)범위가 (left,right)범위와 일치한다면 해당 index 노드가 sum의 답이므로 return tree[index]를 해준다.

3. 만약 일치하지는 않지만 범위가 서로 걸친다면 (left,right)을 반으로 나눈 (left,mid), (mid+1,right) 범위를 가진 자식 노드들을 재귀로 탐색해 나간다.