[python] 세그먼트 트리 구현하기 (segment tree) - 최소값 구하기

세그먼트 트리

https://www.acmicpc.net/blog/view/9

백준님 글

배열 크기 결정

세그먼트 트리를 만들기 위해서 C++의 경우, vector 자료형을 사용하는데 python으로 구현할 것이기 때문에 배열을 사용하려고 한다. 배열의 크기가 얼마나 되어야 할지 가늠하기 위해 크기를 결정해보자.

2의 제곱 수가 리프 노드의 수(n)로 주어진다면 ( 분석해야할 data 갯수가 2의 제곱 수라면 ) 각 단계마다 1/2씩 줄어나가므로 log(n)이 필요한 총 단계(화살표)가 되고 전체 층의 갯수, 즉 높이(h)는 log(n)+1이 된다.

최상위 노드부터 2^(0)개의 노드로 시작하여 층이 증가할 수록 2의 지수 또한 증가하기 때문에 꽉찬 이중 트리(full binary tree)의 총 노드 수는 2^(0)+2^(1)+...+2^(h-1) (h는 높이) 가 되어 2^(h)-1 개 이다.

h = log(n)+1 이므로 필요한 노드의 수는 2^(log(n)+1)-1 개이다.

2의 제곱 수가 아닌 경우도 똑같이 계산할 수 있다. 다만 정수가 나오지 않으므로 올림해주면 된다. 위의 그림에 표시된 회색 동그라미들은 존재하지만 사용하지 않는 배열의 요소들이다.

따라서 처음에 tree 배열의 크기는 2^(log(n)+1)-1이고 다음과 같이 선언할 수 있다.

tree_list = [0]*(pow(2,math.ceil(math.log(len(a),2))+1))-1 #a는 배열
# math.log(갯수,2) => log(갯수)
# math.ceil => 올림
# pow(2,숫자) => 2^(숫자)

일단 필자는 최소의 값을 구하기 위해 세그먼트 트리를 만들고 있으므로 모든 값을 0으로 초기화한다.

init

일단 세그먼트 트리를 주어진 데이터를 이용해서 채워야한다. 따라서 init 함수를 따로 구현하여 세그먼트 트리를 채우기로 하자. init 함수는 2개씩 검사하여 더 작은 수를 위로 올리는 트리의 형태를 띄어야하므로 재귀함수로 구현할 것이다. 이 때, 기저사례는(탈출 조건은) 리프노드이다.

def init(a, tree, node, start, end):
    if start == end:
        tree[node-1] = start
    else:
        init(a,tree,node*2,start,(start+end)//2)
        init(a,tree,node*2+1,(start+end)//2+1,end)
        if a[tree[node*2-1]]<a[tree[node*2]]:
            tree[node-1] = tree[node*2-1]
        else:
            tree[node-1] = tree[node*2]

하나의 예제를 가지고 해당 코드를 따라가며 값들을 좇다보면 어떻게 이런 코드가 구성되었는지 알것이다. 트리에는 주어진 배열의 값이 들어가는 것이 아니라 인덱스가 들어간다. 필자는 일전에 언급한 바와 같이 가장 작은 최솟값을 찾기 위해 세그먼트 트리를 다음과 같이 구성하였다.

query

세그먼트 트리를 이용하여 최솟값을 가져오는 과정에서 배열 전체를 따지자면 배열의 0번째 index 값에 해당하는 부분만 살펴보면 된다. 최상위 노드가 배열의 첫번째 요소로 저장되어있기 때문이다. 하지만 여러 번 다른 범위로 접근하면서 그 범위에서의 최솟값을 찾으려고 할 때는 어떻게 해야할까? 매번 범위에 해당하는 세그먼트 트리를 새로 만들어서 첫번째 요소를 가져오는 것도 한 방법이겠지만 시간이 오래 걸릴 것이다. 따라서 query라는 함수를 만들어 해당 범위에서의 최솟값을 계산하도록 한다.

def query(a,tree,node,start,end,i,j):
    if i>end or j<start: # 주어진 범위가 전체 범위에 완전히 속하지 않는다면
        return -1 
    if i<=start and end<=j: # 주어진 범위가 전체 범위를 완전히 포함한다면
        return tree[node-1] # 해당 범위의 최소값 index를 return 함

    # 주어진 범위가 전체 범위에 속해있다면
    # ex_) 0~4 index의 값이 주어졌을 때, 1~3 index에 대한 최소값을 구하고 싶음
    m1 = query(a,tree,2*node,start,(start+end)//2,i,j) # 전체 범위를 반으로 자른 범위를 넘기고 왼쪽 노드 값을 넘김
    m2 = query(a,tree,2*node+1,(start+end)//2+1,end,i,j) # 전체 범위를 반으로 자른 범위를 넘기고 오른쪽 노드 값을 넘김
    if m1==-1:
        return m2
    elif m2==-1:
        return m1
    else:
        if a[m1]<=a[m2]:
            return m1
        else:
            return m2

해당 그림과 코드를 같이 살펴보면 해당 그림에서 1-4 사이의 최소값을 찾고자 할 때 필요한 노드는

위의 3가지 노드이다. query 함수에 i,j 매개변수를 통해 1,4를 넘기면 처음에는 if문 2개를 통과하고 m1,m2를 계산하게 된다. 이 때 m1은 0-2, m2는 3-4 index의 최소값을 각각 받아오므로 그림의 2번째 층을 의미한다고 볼 수 있다. 따라서 매개변수로 넘겨지는 node 값도 이에 맞춰 바뀌었다.

m1으로 가서 0-2는 또 if문 2개를 통과해서 m1이 0-1, m2가 2로 주어진다.

다시 m1으로 들어가면 m1이 0-0, m2가 1-1로 나뉘어 각각 m1은 범위에 속하지 않으므로 -1, m2는 범위에 속하므로 index 1의 값이 return 된다. (맨 아래쪽 왼쪽의 주황색 1 노드) m1,m2 값이 return 되었으므로 아래의 if 문을 통해서 비교해보면 m1이 -1이므로 m2 값을 return하게 된다.

그렇다면 위에서 m1은 index 1의 값이고, m2는 2가 해당 범위에 속하므로 index 2의 값인데 둘다 -1이 아니므로 if문의 2개의 조건을 통과하고 마지막 else 문에서 1과 2 중 주어진 배열에서 더 작은 수를 가진 index가 return 된다.

여기까지 하면 첫번째 문단에서 m1을 구한 것이다. 이는 1-2 범위의 최소값을 구한 것이고 똑같은 방식으로 m2도 구하면 m2는 query 함수의 2번째 if문에 해당하므로 해당 노드에 있는 index를 바로 반환한다.

이 둘을 비교해서 마지막 최소값을 구하게 된다.