[Algorithm] 힙(Heap) 구현하기 (Swift)

힙(Heap) 이란?

힙은 최댓값 혹은 최솟값을 찾아내는 연산을 빠르게 하기 위해 이진트리를 기본으로 하는 자료구조이다. 구조는 자식노드, 부모노드로 구성이 되어있는데 직관적으로 알 수 있겠지만...(?) 삼각형으로 생각한다면 위 꼭짓점이 부모노드, 아래 두 꼭짓점에 위치한 노드가 자식노드라고 할 수 있다.

 

부모노드의 키값이 자식노드의 키 값보다 항상 큰 힙을 최대 힙, 부모노드의 키 값이 자식노드의 키 값보다 항상 작은 힙을 최소 힙이라고 부른다. 이 두 조건만 만족하게 된다면 숫자의 크기에 따른 순서는 상관이 없다.(예를 들면, 아래의 그림에서 보듯이 20이 느낌상 맨끝에 있어야 할 것 같은데 그렇지 않다던가...) 각 노드 위에 붙은 번호는 index라고 생각하면 된다.

최대 힙 (왼쪽) 최소 힙 (오른쪽)

힙에서의 index

index가 힙을 구현하는데 아주 중요하다!! index를 보면 부모노드가 자식노드 index/2 라는 것을 쉽게 확인할 수 있다. 그리고 자식노드는 index*2 이거나 index*2+1 이 된다.

 

힙에 데이터를 추가하게 된다면 append로 추가하게 될 텐데, 이는 맨 마지막에 추가된다. 그 이후 계속 부모노드와 값을 비교해서 최대힙일 경우 부모노드보다 크면, 최소힙일 경우 부모노드보다 작으면 위로 올리는 과정을 반복해 준다면 힙의 자료구조를 지킬(?) 수 있을 것이다.

 

힙에서 데이터를 추출하게 된다면 index = 1의 데이터를 추출하게 된다. 그리고 맨 마지막 노드를 맨 첫 번째로 보내준 후 양쪽 자식노드들 중 최대힙일 경우 큰 자식노드와 비교를 하고 최소힙일 경우 작은 자식노드와 비교를 하여 힙의 자료구조를 생각하며 쭉쭉 내려주면 된다.

 

슬프게도 Swift에서는 이러한 힙 자료구조를 지원해주지 않고 있다... (ㅠㅠ?) 그래서 직접 구현해봤다.

 

최대 힙 구현

생각보다 간단했다. index에 관한 내용은 사실 힙에 대한 개념이 있다면 따로 함수를 구현할 필요 없이 index자체로 충분히 요리(?)가 가능하고, push, pop부분만 구현해 주면 되었다. 아직 부족해서 일수도 있지만 guard는 직관적으로 반대로 생각해줘야 한다는 느낌 때문에 가독성이 조금 떨어지는 것 같아서 최대한 if를 사용하였다.

struct MaxHeap <T: Comparable> {
    var heap: [T] = []

    func isEmpty() -> Bool {
        heap.count <= 1
    }

    mutating func push(_ data: T) {
        if isEmpty() { heap.append(data) }
        var index = heap.count
        heap.append(data)

        while index > 1 {
            let parent = heap[index/2]
            if parent > data { break }
            heap[index] = parent
            index /= 2
        }
        heap[index] = data
    }

    mutating func pop() -> T? {
        if isEmpty() { return nil }
        let item = heap[1]
        let data = heap.removeLast()
        let size = heap.count - 1

        if isEmpty() { return item }
        heap[1] = data
        
        var (parent, child) = (1, 2)
        while child <= size {
            if child < size && heap[child] < heap[child+1] {
                child += 1
            }
            if heap[parent] > heap[child] { break }
            heap.swapAt(parent, child)
            parent = child
            child = parent * 2
        }
        
        return item
    }
}

 

최소 힙 구현

최대 힙과 개념은 비슷하지만 push, pop 부분에서 조건만 살짝 바꿔주면 바로 구현이 가능했다.

struct MinHeap <T: Comparable> {
    var heap: [T] = []

    func isEmpty() -> Bool {
        heap.count <= 1
    }

    mutating func push(_ data: T) {
        if isEmpty() { heap.append(data) }
        var index = heap.count
        heap.append(data)

        while index > 1 {
            let parent = heap[index/2]
            if parent < data { break }
            heap[index] = parent
            index /= 2
        }
        heap[index] = data
    }

    mutating func pop() -> T? {
        if isEmpty() { return nil }
        let item = heap[1]
        let data = heap.removeLast()
        let size = heap.count - 1
        
        if isEmpty() { return item }
        heap[1] = data
        
        var (parent, child) = (1, 2)
        while child <= size {
            if child < size && heap[child] > heap[child+1] {
                child += 1
            }
            if heap[parent] < heap[child] { break }
            heap.swapAt(parent, child)
            parent = child
            child = parent * 2
        }
        
        return item
    }
}