기본적인 세그먼트 트리 구현

문제: https://www.hackerrank.com/challenges/range-minimum-query

숫자 n개가 주어지고, 일정 구간 내에서 최솟값을 찾는 쿼리 m개에 모두 답해야 한다. 각 숫자는 10^5를 넘지 않는다.

n, m이 최대 10만이기 때문에 항상 구간 내의 모든 값을 확인해서 최솟값을 찾으면 시간복잡도가 O(nm)이여서 통과하지 못한다. 세그먼트 트리를 이용하면 트리 구축에 O(n), 쿼리 하나에 O(logn) 시간이 걸리므로 O(n) + O(mlogn) ⊆ O(mlogn)으로 테스트를 통과할 수 있다.

세그먼트 트리는 크게 세 가지 구현이 가능하다.

1. 트리 기반, 탑 다운 방식

2. 트리 기반, 바텀 업 방식

3. 배열 기반

이름에 트리가 들어가지만 세그먼트 트리는 배열로 구현할 수 있고, 이것이 메모리, 속도 면에서 가장 효율적인 구현이다, 하지만 하스켈에서 이 방식으로 구현하려면 Data.Array를 써야 하는데, ST 모나드를 쓰지 않는 한 Data.Array는 부분 변경이 불가능하고 원소를 하나만 고쳐도 새로운 배열을 생성해야 한다. 그래서 세그먼트 트리의 한 원소를 변경하거나 지연 전파(lazy propagation) 같은 기법을 구현하기에 부적합하다. 이 글에서는 1번과 2번 방식을 구현한다. 3번 방법은 http://codeforces.com/blog/entry/18051를 참조할 것.

세그먼트 트리 자료형

먼저 세그먼트 트리를 나타낼 자료형을 정의한다. 간단한 이진 트리 형식으로 정의할 수 있다.

type NodeData = (Int, Int, Int) -- left bound, right bound, minimum value in the range

data SegTree = Empty | Node NodeData SegTree SegTree deriving Show

rangeMin (Node (_,_,x) _ _) = x

탑 다운 방식

전체 구간을 포함하는 노드로 시작해 구간을 반절로 쪼개며 세그먼트 트리를 구축한다. 각 노드는 구간과 그 구간에서의 최솟값을 보관하고, 두 자식 노드를 가진다. 각 자식 노드는 구간의 왼쪽 반절과 오른쪽 반절에 대한 정보를 가진다. 구간의 길이가 1이 될 때까지 분할을 반복한다.

buildTree :: [Int] -> Int -> SegTree

buildTree nums n = f nums (0,n-1) where

    f [x] (l,r) = Node (l,r,x) Empty Empty

    f xs (l,r) = Node (l,r,minLR) ltree rtree where

        mid = (l + r) `div` 2

        half = mid - l + 1

        ltree = f (take half xs) (l,mid)

        rtree = f (drop half xs) (mid+1,r)

        minLR = min (rangeMin ltree) (rangeMin rtree)

바텀 업 방식

n개 원소 각각에 대응하는 n개 노드를 생성한다. 그리고 2개씩 병합하여 부모 노드를 생성한다. 원소가 홀수 개이면 마지막 원소는 그대로 놔둔다. 새로 생성된 노드들을 다시 2개씩 병합한다. 이 과정을 노드가 하나 남을 때까지 반복한다.

buildTree :: [Int] -> Int -> SegTree

buildTree nums n = f tree0 where

    tree0 = [Node (i,i,x) Empty Empty | (i,x) <- zip [0..n-1] nums]

    collect [] = []

    collect [x] = [x]

    collect (x:y:xs) = merge x y : collect xs

    merge ltree@(Node (l1,r1,x) _ _) rtree@(Node (l2,r2,y) _ _) = Node (l1,r2,min x y) ltree rtree

    f [x] = x

    f xs = f (collect xs)

최솟값 쿼리

어느 방식으로 세그먼트 트리를 구축했던 간에 똑같은 쿼리 함수를 사용할 수 있다. 쿼리 구간이 [L, R]이라 하면 루트 노드부터 검색을 시작한다. 노드가 나타내는 범위 [min, max]가 [L, R]에 포함되면 자신의 최솟값을 반환하고 끝난다. 아예 [L, R]을 벗어나면 유효하지 않은 값을 반환한다. (여기서는 문제가 간단하므로 최솟값이 될 일이 없는 임의의 큰 값을 반환하도록 구현) [min, max]와 [L,R]이 일부만 겹치면 두 자식 노드에 쿼리를 위임한다.

query :: SegTree -> [Int] -> Int

query Empty _ = 10^6 -- maximum element is 10^5

query (Node (rangeL, rangeR, minValue) ltree rtree) range@[l,r]

    | r < rangeL || rangeR < l         = 10^6

    | l <= rangeL && rangeR <= r   = minValue

    | otherwise                             = min retL retR

    where

        retL = query ltree range

        retR = query rtree range

단순히 이 문제를 풀기 위해 작성된 함수이므로 범용적으로 고치려면 10^6 같은 임의의 상수를 쓰지 않아야 한다. 임의의 큰 수 대신 Maybe를 이용하면 안전한 구현이 가능하다.

세그먼트 트리를 구축했으면 남은 일은 쿼리를 읽어서 답하는 것 뿐이다.

트리 시각화

탑 다운 방식과 바텀 업 방식은 최종 트리의 모양이 달라진다. 높이나 쿼리 속도에 큰 영향은 주지 않는다.

탑 다운 방식

바텀 업 방식