세그먼트 트리, python

 

세그먼트 트리란?

 

세그먼트 트리란 완전 이진 트리에 배열의 구간 정보(최대, 최솟값, 합, 곱 등)를 저장하고 빠르게 구간 쿼리를 처리할 수 있는 자료구조이다.

완전 이진트리란 리프노드를 제외한 모든 노드들의 자식노드들을 2개씩 가지고 있어야 하고 왼쪽부터 채워져 있어야 한다.

 

이 트리는 리프노드 8,9,10,11,12를 제외한 모든 노드들의 자식노드가 2개씩 있고 왼쪽부터 채워져 있으므로 이진트리이다.

 

 

완전 이진트리는 배열로 구현할 수 있다.

루트노드의 인덱스는 1번이다.

왼쪽 자식노드는 부모노드의 인덱스 * 2번이 되고, 오른쪽 자식노드는 부모노드의 인덱스 * 2  + 1이 된다.

 

 

 

 

 

 

 

 

 

 

세그먼트 트리의 크기 구하기

 

 

세그먼트 트리는 완전 이진트리로 구현할 수 있다.

배열이 [6 , 3 , 7, 2, 4 , 1]이라고 하면 먼저 세그먼트 트리의 리프노드에 각 배열의 요소를 저장한다.

 

따라서 세그먼트 트리의 리프노드의 개수는 배열의 요소가 전부 들어가야 하므로 배열의 길이보다 크면서 2^n을 만족해야 한다.

리프노드의 개수가 2^n개 일 때, 전체 노드의 개수는 1+2+2^2+...+2^(n-1) 이므로 2^(n+1)-1 개가 된다.

따라서 세그먼트 트리를 구현할 배열의 길이를 2^(n+1)로 해주어야 한다.

배열의 길이보다 크면서 2^n을 만족하는 수를 size라고 하면 2*size가 된다.

 

 

 

예시에서는 배열의 크기가 6이므로 6보다 크면서 2^n을 만족하는 수는 8이다.

그러므로 세그먼트 트리 배열의 길이는 16이 된다.

 

 

 

 

 

 

세그먼트 트리 초기값 설정하기

 

세그먼트 트리의 길이를 구했으면 이제 각 노드에 배열 구간의 정보를 업데이트해줘야 한다.

 

먼저 배열을 구간의 크기를 1로 쪼개고 배열의 각 값을 리프노드에 넣어준다.

그림의 파란색 부분에 차례대로 배열의 요소가 들어가게 된다.

이진트리에서 마지막 레벨의 가장 왼쪽 노드의 인덱스는 size가 된다.

따라서 segtree [size + i] = arr [i]를 넣어준다.

 

 

 

 

이제 왼쪽 자식노드와 오른쪽 자식노드의 값을 가지고 부모노드의 값을 업데이트해준다.

만약 구간의 최솟값을 구하고 싶다면 부모노드에는 왼쪽 자식노드의 값과 오른쪽 자식노드의 값 중 작은 값을 넣어주면 된다.

구간의 합을 구하고 싶다면 왼쪽 자식노드와 오른쪽 자식노드의 합을 넣어주면 된다.

 

 

 

 

 

 

 

이렇게 루트노드까지 갱신했다면, 세그먼트 트리의 초기 세팅이 끝난 것이다.

 

 

 

 

 

 

 

 

 

세그먼트 트리를 사용한 구간 쿼리

 

이제 이 트리를 사용해서 배열의 (a, b) 구간의 정보를 구해야 한다.

 

 

 

부모노드의 범위는 왼쪽 자식노드와 오른쪽 자식노드의 범위를 합친 것이다.

그리고 완전 이진트리에서 오른쪽 자식노드는 짝수이고 왼쪽 자식노드는 홀수이다.

이 말은 a의 부모노드의 범위가 a 이상에 포함이 되려면, a는 짝수여야 한다.

마찬가지로 b의 부모노드의 범위가 b 이하에 포함되려면 b는 홀수여야 한다.

 

 

 

a가 짝수인 경우 부모노드로 올려준다. 이렇게 되면 a이상이면서 미리 정보를 계산해 놓은 구간을 건너뛸 수 있게 된다.

마찬가지로 b가 홀수인 경우도 부모노드로 올려주게 되면 b 이하면서 미리 정보를 계산해 놓은 구간만큼 건너뛸 수 있다,

 

 

a가 홀수, 또는 b가 짝수인 경우 해당 노드의 값을 최종 결과에 먼저 반영을 한 뒤, a를 한 칸 오른쪽으로 옮겨 짝수로 만들고, b는 한 칸 왼쪽으로 옮겨 홀수로 만들면서 범위를 안쪽으로 좁혀 나가면서 a와 b가 교차되면 종료한다.

 

 

 

 

 

위에서 배열 [6 , 3 , 7, 2, 4 ,1]로 생성된 세그먼트 트리의 (1,4) 구간 합을 구하는 과정이다. 

 

 

세그먼트 트리를 초기화하면 위와 같이 된다.

 

 

 

 

 

 

 

 

 

먼저 (1,4)를 리프노드부터 시작한다

size+1 = 9번, size+4 = 12번 에서 시작한다.

 

 

 

 

 

 

 

 

left는 9이고 홀수이므로 오른쪽 자식노드이다. 따라서 결과값에 합 3을 더해주고 한 칸 오른쪽으로 옮겨주었다.

right는 12고 짝수이므로 왼쪽 자식 노드이다. 결과값에 합 4를 더해주고 한 칸 왼쪽으로 옮겨주었다.

현재 총합은 7이다.

 

 

 

 

 

 

 

left는 10이고 짝수, right는 11이고 홀수이므로 부모노드로 옮겨준다.

 

 

 

 

 

 

left는 5이고 홀수이므로 한 칸 왼쪽으로 옮겨주었고 결과값에 9를 더한다.

right는 짝수이므로 부모노드로 옮겨주었다.

현재 총합은 16이다.

여기서 left가 right보다 커지게 되어 교차되게 된다. 따라서 종료한다.

 

 

 

 

 

 

 

배열의 값이 변경되었을 때

 

만약 배열의 값이 변경되는 경우 세그먼트 트리도 업데이트해주어야 한다.

먼저 배열의 변경된 인덱스에 해당하는 리프노드를 변경된 값으로 교체해 주고, 부모노드로 올라가면서 부모노드가 저장하고 있는 구간의 정보를 다시 계산하여 업데이트해주면 된다.

 

 

 

[6 , 3 , 7, 2, 4 ,1] 배열의 4번째 요소를 2에서 10으로 바꿔주었다. 먼저 리프노드의 값을 바꿔준다.

 

 

 

 

 

 

완전 이진트리에서 부모노드의 인덱스는 현재노드를 2로 나눈 몫이다.

부모노드로 이동하고 기존 2에서 10이 되었으므로 차이값인 8을 더해준다

 

 

 

 

 

 

 

부모노드로 올라가면서 루트노드까지 차이값을 계속 업데이트해준다.

 

 

 

 

 

python 구현

 

 

# 배열의 길이보다 크면서 2^n값을 찾음
size = 1
while size < n:
    size *= 2

# 세그먼트 트리 생성
segtree = [0]*(size*2)

# 리프노드에 왼쪽부터 차례대로 배열의 각 요소를 넣음
for i in range(n):
    segtree[size+i] = arr[i]

# 합 세그먼트 트리이므로 부모노드는 왼쪽 자식노드 + 오른쪽 자식노드의 합
# 최솟값 세그먼트 트리면 왼쪽,오른쪽 자식노드의 최솟값을 저장
# 곱 세그먼트 트리면 왼쪽, 오른쪽 자식노드의 곱을 저장
for i in range(size-1,0,-1):
    segtree[i] = segtree[i*2] + segtree[i*2+1]

 

먼저 세그먼트 트리의 크기를 구하고 초기값을 넣어준다.

 

 

 

# 구간합을 구하는 쿼리
def query(left,right):
    result = 0
    # 리프노드부터 시작
    left += size
    right += size
    
    # left, right가 교차할때까지 반복
    while left<=right:
    	# left가 홀수라면 결과에 반영하고 한칸 오른쪽으로 옮김
        if left % 2 == 1:
            # 합 세그먼트 트리 이므로 결과값에 더해줌
            result += segtree[left]
            left += 1
        
        # right가 짝수라면 결과에 반영하고 한칸 왼쪽으로 옮김
        if right % 2 == 0:
            result += segtree[right]
            right -= 1
            
        # 부모노드로 옮김
        left //= 2
        right //= 2
        
    return result

 

left와 right를 조절하면서 쿼리를 수행한다.

 

 

 

def update(i,value):
    diff = value - arr[i]
    
    # 리프노드부터 시작
    i += size
    
    # 루트노드일때까지 반복
    while i>0:
        # 합 세그먼트 트리이므로 차이만큼 더해줌
        segtree[i] = segtree[i] + diff
        
        # 부모노드로 옮김
        i //= 2

 

리프노드에서 시작해 부모노드로 옮겨가면서 차이값을 업데이트해준다

 

 

 

 

 

세그먼트 트리를 사용해서 해결할 수 있는 문제

 

[백준] 2357 최솟값과 최댓값

[백준] 2357 최솟값과 최댓값

 

[백준] 2042 구간 합 구하기

[백준] 2042 구간 합 구하기

 

[백준] 7578 공장

[백준] 7578 공장

 

[백준] 1517 버블 소트

[백준] 1517 버블 소트