[자료구조] Segment Tree : 구간 연산 구간 변경

백준님의 블로그내용으로 공부한 것을 정리한 포스팅입니다.

---

세그먼트 트리.(Segment Tree)

Intro

1. 구간 l,r 이 주어졋을 때 A[l]~A[r]까지의 합
2. i 번재 수를 v로 바꾸기 A[i] = v

수행해야 하는 연산은 최대 M번.

1번연산하는데 O(N) 2번연산 O(1) 총 시간복잡고 O(MN)

2번 연산이 없다고 가정해봅시다.

수를 바꾸는 경우가 없기 때문에 합도 변하지 않습니다.

따라서 앞에서부터 차례대로 합을 구해놓는 방식으로 문제를 풀 수 있습니다.

S[i] = A[1]+...+A[i]

i~j까지 합은 S[j]-S[i-1]

S[0] = A[0]
for i in range(N):
    S[i] = S[i-1] + A[i]

여기서 2번연산을 하려면 수가 바뀔때마다 S를 변경해줘야 합니다. 가장 앞에있는 0번째 수가 바뀌었을 경우 모든 S배열을 변경해야 하기 때문에 시간 복잡도는 O(N)이 걸리게 됩니다. 따라서 M과 N이 매우 큰 경우에는 시간이 너무 오래걸리게 됩니다.

---

세그먼트트리

를 이용하면 1번연산을 O(logN) 2번연산을 O(logN)만에 수행할 수 있습니다.

세그먼트 트리의 리프 노드와 리프노드가 아닌 다른 노드는 다음과 같은 의미를 가집니다.

• 리프 노드 : 배열의 그 수 자체
• 다른 노드 : 왼쪽 자식과 오른쪽 자식의 합을 저장함.

어떤 노드 번호가 x일때 왼쪽 자식의 번호는 2x 오른쪽 자식의 번호는 2x+1이 됩니다.

N=10인 겨우 세그먼트 트리는 다음과 같습니다.

위의 그림은 각 노드가 저장하고 있는 합의 범위를 나타낸 그림입니다. 여기에 각 노드의 번호를 그림으로 나타내면 다음과 같습니다.

트리만들기

만약! N이 2의 제곱꼴인 경우에는 Full Binary Tree입니다. 또 그 때 높이는 logN입니다. 리프 노드가 N개인 Full Binary Tree는 필요한 노드의 개수가 2N-1입니다.

N의 제곱꼴이 아닌경우는 높이가 H=[logN]이고, 총 세그먼트 트리를 만드는데 필요한 배열의 크기는 2^(H+1)-1개가 됩니다.

다음과 같은 과정을 거쳐서 Segment Tree를 만들 수 있습니다.

C++의 경우vector자료형을 쓰는데 파이썬은 배열을 사용해야한다.

# 리스트 a, 세그트리, 노드번호, 노드가 담당하는 합 S의 범위 start~end
def segtree(list_a,tree,node,start,end):
    if start==end:
        return tree[node]=a[start]
    else:
        return tree[node] = segtree(a,tree,node*2,start,(start+end)/2) + segtree(a,tree,node*2+1,(start+end)/2,end)

• start==end인 경우는 node가 리프 노드인 경우입니다. 리프 노드는 배열의 그 원소를 가져야 하기 때문에 tree[node] = a[start] 가 됩니다.
• node의 왼쪽 자식은 node2, 오른쪽 자식은 node2+1이 됩니다. 또 노드가 담당하는 구간이 start,end라면 왼쪽자식은 start,(start+end)/2 오른쪽 자식은 (start+end)/2,end를 담당해야합니다.
• 따라서 재귀 함수를 이용해서 왼쪽 자식과 오른쪽 자식 트리를 만들고 그 합을 저장해야 합니다.

합 찾기

구간 left, right가 주어졌을 때 합을 찾으려면 루트부터 트리를 순회하면서 각 노드가 담당하는 구간과 left, right사이의 관계를 살펴봐야 합니다. 예를들어 0~9까지합을 구하는 경우 루트노트 하나만으로 합을 알 수 있습니다.

2~4까지 합을 구하는 경우는 다음과 같습니다.

5~8

3~9

노드가 담당하고 있는 구간이 start,end고 합을 구해야하는 구간이 left,right면 다음과 같은 4가지 경우로 나누어 질 수 있습니다.

1. left,right와 start,end가 겹치지 않는경우
2. left,right가 start,end를 완전히 포함하는 경우
3. start,end가 left,right를 완전히 포함하는 경우
4. left,right와 start,end가 겹쳐져 있는 경우.

1번 경우에는 if left>end or right <start로 나타낼 수 잇습니다. 이경우는 더이상 탐색을 이어나갈 필요가 없습니다. 따라서 0을 리턴해 탐색을 종료합니다.

2번 경우에는 if left<=start and end<=right 로 나타낼 수 있습니다. 이경우도 더이상 탐색을 이어나갈 필요 없습니다. 구해야 하는 범위가 그 노드범위에 포함되고 그 노드의 자식도 모두 포함되기 때문입니다. 때문에 tree[node]를 리턴해 값을 주고 탐색을 종료합니다.

3번과 4번의 경우에는 왼쪽 자식과 오른조 ㄱ자식을 루트로 하는 트리에서 다시 탐색을 시작해야합니다.

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

수의 변경

중간에 어떤 수를 변경한다면 그 숫자가 포함된 구간을 담당하는 노드를 모두 변경해줘야합니다.

3번째 수를 변경해야 할때 구간을 나타냅니다.

5를 변경

index 번째 수를 val로 변경한다면 그 수가 얼마만큼 변했는지를 알아야 합니다. 이수를 diff 라 하면 diff = val-a[index]로 구할 수 있습니다.

수 변경은 2가지 경우가 있습니다.

1. start,end에 index가 포함되는 경우
2. start,end에 index가 포함되지 않는 경우

node구간에 포함되는 경우에는 diff만큼 증가시켜 합을 변경해 줄 수 있씁니다. tree[node] = tree[node] + diff(차이만큼만 반영)

포함되지 않는경우는 그 자식도 포함되지 않기때문에 탐색을 중단시켜줍니다.

def update(tree,node,start,end,index,diff):
    if index<start or index > end:
        return
    tree[node] = tree[node] + diff
    if start != end:
        update(tree,node*2, start, (start+end)/2, index,diff)
        update(tree,node*2+1,(start+end)/2,end,index,diff)

리프노드가 아닌 경우에는 자식도 변경해 줘야 하기 때문에 start !=end로 리프노드인지 검사해야합니다.

세그먼트 트리를 이용해서 2042번 문제: 구간 합 구하기 번을 풀 수 있습니다. 최소값, 최대값도 합을 구하는 세그먼트 트리를 응용해서 구현할 수 있습니다. 10868번 문제: 최솟값, 2357번 문제: 최솟값과 최댓값

구간 곱 등 구간에 대한 변경과 연산이 있을 시 변형하기 좋다.

import sys
input = sys.stdin.readline

# 세그먼트 트리 생성
def init(node, start, end): 
    # node가 leaf 노드인 경우 배열의 원소 값을 반환.
    if start == end :
        tree[node] = l[start]
        return tree[node]
    else :
        # 재귀함수를 이용하여 왼쪽 자식과 오른쪽 자식 트리를 만들고 합을 저장.
        tree[node] = init(node*2, start, (start+end)//2) + init(node*2+1, (start+end)//2+1, end)
        return tree[node]

# 구간 합 구하기
# node가 담당하는 구간 [start, end]
# 합을 구해야하는 구간 [left, right]
def subSum(node, start, end, left, right) :
    # 겹치지 않기 때문에, 더 이상 탐색을 이어갈 필요가 없다.    
    if left > end or right < start :
        return 0
    # 구해야하는 합의 범위는 [left, right]인데, [start, end]는 그 범위에 모두 포함되고
    # 그 node의 자식도 모두 포함되기 때문에 더 이상 호출을 하는 것은 비효율적이다.
    if left <= start and end <= right :
        return tree[node]
    # 왼쪽 자식과 오른쪽 자식을 루트로 하는 트리에서 다시 탐색을 시작해야한다.
    return subSum(node*2, start, (start+end)//2, left, right) + subSum(node*2 + 1, (start+end)//2+1, end, left, right)


def update(node, start, end, index, diff) :
    if index < start or index > end :
        return
    tree[node] += diff
    # 리프 노드가 아닌 경우에는 자식도 변경해줘야 하기 때문에 검사해야함.
    if start != end :
        update(node*2, start, (start+end)//2, index, diff)
        update(node*2+1, (start+end)//2+1, end, index, diff)


n, m, k = map(int, input().rstrip().split())
l = []
tree = [0] * 3000000

for _ in range(n) :
    l.append(int(input().rstrip()))

init(1, 0, n-1)

for _ in range(m+k) :
    a, b, c = map(int, input().rstrip().split())

    # 1인경우 바꿔주고
    if a == 1 :
        b = b-1
        diff = c - l[b]
        l[b] = c
        update(1, 0, n-1, b, diff)
    # 2인경우 구간합을 반환한다.
    elif a == 2 :                
        print(subSum(1, 0, n-1 ,b-1, c-1))