[Swift] 힙(Heap) 구현하기

배경

파이썬에서는 heap 관련 라이브러리를 제공하기 때문에 heap에 대한 원리만 알아도 문제 푸는데 어려움은 없었다. 하지만 swift에서는 heap 관련 라이브러리가 없어서 사용하려면 내가 직접 구현해야 한다. 따라서 이번 시간에는 최소힙을 구현해보도록 하자(최소힙 구현할 줄 알면 최대힙은 껌).

구현

우선 최소힙의 동작 방식을 살펴보자. 플로우를 내 방식대로 나열해 보겠음.
1. Heap Object를 생성한다.
2. 정수 값을 넣는다. -> 힙을 채우는 과정
3. 정수 값을 추출한다. -> 힙에서 최소 값을 꺼내는 과정

사용하는 방법은 어렵지 않으나 구현하려면 살짝 머리를 써야 한다. 우선 Heap 클래스를 하나 생성해보자.

class BinaryHeap {
    
    var items = [Int]()
}

items라는 이름의 배열을 하나 생성하였다. 여기에 정수값을 넣어주게 된다. 무턱대고 넣으면 힙이 아니니 어떻게 넣어야할까?

우선 계산하기 쉽게 items의 첫번째 인덱스는 사용하지 않는다. 따라서 인덱스 1부터 시작할 것이다. 고로 힙을 초기화할 때 초기화할 값을 두 번 넣어주겠다.

class BinaryHeap {
    
    var items = [Int]()
    
    init(_ val: Int) {
        items.append(val)
        items.append(val)
    }
}

삽입

이제 값을 넣어줄 차례인데 무조건 인덱스 1에는 가장 작은 값이 온다. 따라서 값을 넣어줄 때 자신의 부모와 비교하면서 위로 올라가는 과정이 필요하다!
위의 그림을 보면 1번 인덱스가 가장 작은 값이 들어 있는 것을 확인할 수 있다. 트리와 배열의 각 값의 인덱스를 비교해보면 값을 어떤 순서로 비교해 나가야할지 알 수 있다.

이제 items에 값을 넣는 코드를 보면서 이해해보자.

func percolate_up() {
        var i = items.count-1
        var parent = Int(i / 2)
        while parent > 0 {
            if items[i] < items[parent] {
                items.swapAt(i, parent)
            }
            i = parent
            parent = Int(i / 2)
        }
    }
    
    func insert(k: Int) {
        items.append(k)
        percolate_up()
    }

우선 값을 넣으려면 insert(k:)을 사용하게 되는데 우선 값을 items에 넣고 percolate_up()을 통해 자신의 부모와 비교해가며 적절한 위치에 값을 넣어주는 작업을 거치게 된다. 참고로 percolate는 '스며들다'라는 뜻인데 up이니까 '위로 스며들다', 즉 위로 비교해가면서 값을 변경해준다는 뜻으로 이해하면 적절할 것 같다.
insert(k: 13)인 경우를 생각해봤을 때 부모 값인 10과 비교를 거치게 된다. 하지만 13은 10보다 크니까 바뀌지 않고 그 자리에 놓이게 된다.
insert(k: 7)을 생각해보면 7은 부모 값인 10보다 작으니 부모 값과 스왑을 거치게 된다.
그 다음 다시 부모 값인 6과 비교하게 되는데 6보다는 크니까 그 자리에 위치하게 된다.

추출

추출하는 과정 또한 어렵지 않다. items 배열에 예쁘게 넣어줬으니 제일 작은 값의 인덱스인 1번 인덱스를 추출하고 그 자리에 제일 작은 값을 채워주면 된다.

    func percolate_down(idx: Int) {
        let left = idx * 2
        let right = idx * 2 + 1
        var smallest = idx
        
        if left <= items.count-1 && items[left] < items[smallest] {
            smallest = left
        }
        
        if right <= items.count-1 && items[right] < items[smallest] {
            smallest = right
        }
        
        if smallest != idx {
            items.swapAt(idx, smallest)
            percolate_down(idx: smallest)
        }
    }
    
    func extract() -> Int {
        let extracted = items[1]
        items[1] = items[items.count-1]
        items.popLast()
        percolate_down(idx: 1)
        return extracted
    }

extract()를 살펴보자. items의 첫번째 인덱스의 값이 제일 작으므로 값을 추출한 다음 items의 맨 마지막 값을 첫번째 인덱스에 넣어준다. 그리고 배열의 마지막 값을 없애준다. 지금까지의 과정을 그림으로 보면 다음과 같다.
그 다음 percolate_down(idx:)를 통해 1번 인덱스에 있는 값을 제 자리를 찾아가도록 해준다. 이 과정에서 1번 인덱스에는 제일 작은 값이 들어가게 된다.

percolate_down(idx:) 코드를 살펴보겠다. left와 right는 각각 1번 인덱스의 왼쪽, 오른쪽 자식 노드를 의미한다. 왼쪽 노드의 값과 비교하여 더 크면 smallest의 값을 left로 바꿔주고, 오른쪽 노드의 값과 비교하여 더 크면 smallest의 값을 right로 바꿔준다.
만약 부모 노드의 값이 더 큰 경우가 있다면 그 다음 smallest와 idx 위치의 값을 swap해주고 percolate_down을 반복해준다.

전체 코드

class BinaryHeap {
    
    var items = [Int]()
    
    init(_ val: Int) {
        items.append(val)
        items.append(val)
    }
    
    func percolate_up() {
        var i = items.count-1
        var parent = Int(i / 2)
        while parent > 0 {
            if items[i] < items[parent] {
                items.swapAt(i, parent)
            }
            i = parent
            parent = Int(i / 2)
        }
    }
    
    func insert(k: Int) {
        items.append(k)
        percolate_up()
    }
    
    func percolate_down(idx: Int) {
        let left = idx * 2
        let right = idx * 2 + 1
        var smallest = idx
        
        if left <= items.count-1 && items[left] < items[smallest] {
            smallest = left
        }
        
        if right <= items.count-1 && items[right] < items[smallest] {
            smallest = right
        }
        
        if smallest != idx {
            items.swapAt(idx, smallest)
            percolate_down(idx: smallest)
        }
    }
    
    func extract() -> Int {
        let extracted = items[1]
        items[1] = items[items.count-1]
        items.popLast()
        percolate_down(idx: 1)
        return extracted
    }
}

참고

https://babbab2.tistory.com/109
https://www.cs.usfca.edu/~galles/visualization/Heap.html