플레찍었으니 정리해보는 알고리즘 노트 4 - 세그먼트 트리

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

2042번: 구간 합 구하기

세그먼트 트리란?

 

여러개의 데이터 중에서 특정 구간의 합(곱, 최대, 최솟값 등)을 빠르게 구할 수 있는 자료구조이다.

투 포인터 역시 구간의 합을 구하는 알고리즘이지만 시간복잡도가 O(N) 이므로

똑같은 데이터가 주어지고 탐색을 여러번 반복해야 하는 상황이라면 효율적이지 못하다.

 

세그먼트 트리는 이진트리 구조를 활용하며 루트 노드(1번 노드) 에는 모든 구간의 합

그 밑의 자식노드(2번 노드, 3번 노드)는 부모 노드의 절반 구간까지의 합을 삽입하는식으로 계속 진행한다.

 

재귀함수를 이용하면 구현이 간단하며 루트 노드를 1번 노드로 설정하면 부모 노드 * 2가 왼쪽 자식 노드로 이어진다.

노드의 총 개수는 데이터의 개수 * 4로 설정하면 충분하다.

import sys
n,m,k = map(int,sys.stdin.readline().rstrip().split())
lst = [int(sys.stdin.readline().strip()) for _ in range(n)]
tree = [0]*n*4

def init(index,start,last): #범위를 좁혀나가는 식으로 재귀함수 이용 구현
    if start >= last:
        tree[index] = lst[start-1]
    else:
        c = int((start+last)/2)
        tree[index] = init(index*2,start,c) + init(index*2+1,c+1,last)
    return tree[index]

def find(index,start,last,left,right):
    if start > right or last < left:
        return 0
    elif start <= left and last >= right:
        return tree[index]
    else:
        c = int((left+right)/2)
        return find(index*2,start,last,left,c) + find(index*2+1,start,last,c+1,right)

def update(d,x,index,start,last):
    if x >= start and x <= last:
        tree[index] += d
        if start != last:
            c = int((start+last)/2)
            update(d,x,index*2,start,c)
            update(d,x,index*2+1,c+1,last)
        
init(1,1,n)

for i in range(m+k):
    a,b,c = map(int,sys.stdin.readline().strip().split())
    
    if a == 1:
        d = c - lst[b-1]
        lst[b-1] = c
        update(d,b,1,1,n)
    else:
        if b > c:
            b,c = c,b
        print(find(1,b,c,1,n))

재귀함수를 이용하여 start,last를 이용 start와 last 사이의 중간값 c를 구한 후

init(index,start,last) = init(index*2,start,c) + init(index*2 + 1,c+1,last)와 같은 식으로 재귀적으로 구현하면 된다.

init 함수와 find 함수를 조금 수정하면 구간의 합 뿐 아니라 구간의 최대값,최솟값, 곱 등도 구할 수 있다.