백준 2042번 구간 합 구하기

N개의 수가 주어지고 연속된 어떤 부분의 합을 구하는 문제입니다. 중간에 i번째 수를 X로 바꾸어주기도 합니다.

이번 문제의 풀이에 앞서 이번 문제의 알고리즘 분류인 세그먼트 트리에 대해 알아보았습니다.

 

 

세그먼트 트리

세그먼트 트리는 여러 개의 연속된 데이터가 존재 할 경우 연속된 어떤 부분의 데이터 합을 구하는 방법에 관한 알고리즘입니다.

 

N개의 변 할 수 있는 데이터에서 연속된 수 M개의 합을 여러번 구하는 방법은 무조건 M번의 반복을 통해서 구할 수 있습니다. N과 M이 엄청 큰 수라면 오랜 시간이 걸리게 될 것입니다.

 

세그먼트 트리는 트리의 서브트리에 포함된 연속된 범위의 데이터들의 합을 각 서브트리의 루트 노드에 저장 할 수 있도록 해줌으로써 좀 더 빠르게 특정 범위의 데이터의 합을 구 할 수 있습니다.

 

데이터의 수정과 특정 범위의 합을 구하는 시간 복잡도가 모두 logN이므로 훨씬 빠른 속도로 답을 구할 수 있습니다.

아래와 같은 연속된 데이터를 예시로 들어보겠습니다.

1 2 3 4 5

세그먼트 트리는 각 서브트리의 루트 노드가 해당 범위의 숫자들의 합을 저장 할 수 있어야 합니다. 이러한 구조가 된다면 리프노드에는 데이터들이 존재하게 됩니다. 그러므로 루트노드에서 리프노드로 내려간 뒤 재귀적으로 다시 올라가면서 각 범위의 합을 서브트리에 저장할 수 있도록 해줍니다.

def recursion(start,end,node):
    if start==end:
        segTree[node]=arr[start]
    else:
        mid=(start+end)//2
        segTree[node]=recursion(start,mid,node*2)+recursion(mid+1,end,node*2+1)
    return segTree[node]

start와 end는 연속된 데이터의 처음과 끝을 지정해줍니다. recursion 함수 실행 시, start=0,end=N-1로 지정해주었습니다.

node의 경우 start~end 범위의 합을 저장할 노드를 뜻합니다. recursion 함수 실행 시, 루트 노드인 1로 지정해주었습니다.

이를 통해 만들어진 트리의 구조는 다음과 같습니다.

세그먼트 트리 1

위의 recursion 함수를 통해 만들어진 세그먼트 트리입니다.

리프 노드에는 연속된 데이터들이 존재하게 되고 각 서브트리의 루트 노드는 자신의 서브트리에 존재하는 리프노드들의 합을 가지고 있습니다.

 

이러한 세그먼트 트리를 배열의 형태로 저장해두고 이를 활용해 특정 구간의 합을 구해줍니다.

0 15 6 9 3 3 4 5 1 2

위의 세그먼트 트리를 배열로 저장했을 때의 모습입니다. 루트 노드의 경우 1에 저장하는 것이 로직 상 맞기 때문에 0번 인덱스에 0을 채워두었습니다.

 

위의 세그먼트 트리로 2~5 구간의 합을 구할 경우

세그먼트 트리 2

연속된 데이터 2 3 4 5를 구해보도록 하겠습니다.

2와 3은 리프노드까지 내려가야지만 해당 원소들의 합을 구할 수 있습니다.

하지만, 4와 5의 경우 9를 루트노드로 하는 서브트리의 원소들이 모두 범위에 포함되어 있으므로 해당 루트노드만 확인함으로써 범위에 포함된 원소들의 합을 바로 찾아낼 수 있습니다.

이처럼 빠르게 특정 구간의 합을 찾아 낼 수 있도록 해주는 알고리즘이 세그먼트 트리 알고리즘입니다.

 

 

숫자 교체

문제에서 특정 i번째 데이터를 X로 바꾸는 경우도 존재했습니다.

i번째 데이터를 교체 할 경우 해당 원소를 포함하는 서브트리의 루트노드들도 바뀌어야 하기 때문에 재귀를 통해 해결했습니다.

def change(start,end,node,target,num):
    if start==end:
        segTree[node]=num
        return segTree[node]

    mid=(start+end)//2

    if start<=target<=mid:
        segTree[node]=change(start,mid,node*2,target,num)+(segTree[node*2+1] if node*2+1<N*4 else 0)
    elif mid+1<=target<=end:
        segTree[node]=change(mid+1,end,node*2+1,target,num)+(segTree[node*2] if node*2<N*4 else 0)

    return segTree[node]

start, end는 연속된 데이터의 처음과 끝으로 0,N-1을 지정해주었습니다.

node는 start~end 범위의 데이터 합을 가지고 있는 노드로 처음에는 루트 노드를 가리키므로 1을 지정해주었습니다.

target의 경우 바뀌게 되는 i번째 노드의 i를 가리킵니다. num은 바뀌게 될 숫자입니다.

 

start==end일 경우는 리프노드까지 내려온 경우이므로 해당 원소를 num으로 바꿔주고 리턴해줍니다.

start!=end일 경우 target이 포함된 범위쪽으로 내려갈 수 있도록 해줍니다. target이 포함된 범위라면 해당 범위의 합이 바뀌게 됨으로 재귀 호출로 반환 받은 값을 통해 해당 범위의 서브트리 노드의 값 또한 교체해줍니다.

 

 

범위의 합 구하기

앞에서 살펴 보았듯이 주어진 범위의 합을 구하기 위해 서브트리의 루트 노드에서 리프노드로 내려가지 않을 경우도 있고 리프노드까지 내려가는 경우도 있었습니다. 저와 같은 경우 합을 찾아야 하는 범위를 targetStart, targetEnd로 지정해주고 현재 지정된 범위를 start,end로 지정해 구분해주었습니다.

 

이 start~end 범위가 targetStart~targetEnd 범위에 포함되었을 경우 합을 구하려는 범위에 현재 범위가 포함되었다는 뜻이므로 더 이상 내려가지 않고 현재 노드의 값을 반환해줍니다.

targetStart~targetEnd 범위가 start~end범위에 걸쳐진 경우 걸쳐진 범위의 합만 구해줄 수 있도록 했습니다.

def findSum(targetStart,targetEnd,start,end,node):
    if start==end or (start>=targetStart and end<=targetEnd):
        return segTree[node]

    returnValue=0
    mid=(start+end)//2

    if targetStart<=mid:
        returnValue+=findSum(targetStart,targetEnd,start,mid,node*2)
    if mid+1<=targetEnd:
        returnValue+=findSum(targetStart,targetEnd,mid+1,end,node*2+1)

    return returnValue

가장 위쪽에서 범위에 포함된 원소이거나 현재 범위가 target 범위에 포함된 경우에 해당 노드의 값을 반환해줍니다.

target 범위가 현재 범위에 걸쳐진 경우 걸쳐진 범위의 합만을 구할 수 있도록 해줍니다.

 

 

 

 

해결

세그먼트 트리의 문제라는 것을 알고 해당 개념을 익힌 뒤 풀이하여 한 번에 맞출 수 있었던 문제입니다. 세그먼트 트리 관련된 문제들이 모두 백준에서 높은 난이도를 가지고 있는 것을 확인했는데 종종 풀어주면 좋을 것 같습니다.