[Swift] 힙(Heap) 구현하기

Heap은 트리 기반의 자료구조로, 특정 우선순위에 따라 요소들을 정렬하는 기능을 제공합니다.

Heap은 크게 두 가지 유형이 있습니다.

최소 힙은 작은 값이 먼저 나오도록 정렬하고, 최대 힙은 큰 값이 먼저 나오도록 정렬합니다.

 

제네릭을 이용해 Heap 구조체를 정의하고 주요 기능을 구현해보겠습니다.

struct Heap<T> {
    private var elements: [T] = []
    private let priorityFunction: (T, T) -> Bool
    
    var isEmpty: Bool {
        return elements.isEmpty
    }
    
    var count: Int {
        return elements.count
    }
    
    var peek: T? {
        return elements.first
    }
    
    init(priorityFunction: @escaping (T, T) -> Bool) {
        self.priorityFunction = priorityFunction
    }
    
    mutating func push(_ element: T) {
        elements.append(element)
        siftUp(from: elements.count - 1)
    }
    
    mutating func pop() -> T? {
        guard !isEmpty else { return nil }
        elements.swapAt(0, count - 1)
        let element = elements.removeLast()
        if !isEmpty {
            siftDown(from: 0)
        }
        return element
    }
    
    private mutating func siftUp(from index: Int) {
        var childIndex = index
        let child = elements[childIndex]
        var parentIndex = self.parentIndex(of: childIndex)
        
        while childIndex > 0 && priorityFunction(child, elements[parentIndex]) {
            elements[childIndex] = elements[parentIndex]
            childIndex = parentIndex
            parentIndex = self.parentIndex(of: childIndex)
        }
        
        elements[childIndex] = child
    }
    
    private mutating func siftDown(from index: Int) {
        let parent = elements[index]
        var parentIndex = index
        
        while true {
            let leftChildIndex = self.leftChildIndex(of: parentIndex)
            let rightChildIndex = leftChildIndex + 1
            
            var candidateIndex = parentIndex
            
            if leftChildIndex < count && priorityFunction(elements[leftChildIndex], elements[candidateIndex]) {
                candidateIndex = leftChildIndex
            }
            
            if rightChildIndex < count && priorityFunction(elements[rightChildIndex], elements[candidateIndex]) {
                candidateIndex = rightChildIndex
            }
            
            if candidateIndex == parentIndex {
                return
            }
            
            elements.swapAt(parentIndex, candidateIndex)
            parentIndex = candidateIndex
        }
        elements[parentIndex] = parent
    }
    
    private func parentIndex(of index: Int) -> Int {
        return (index - 1) / 2
    }
    
    private func leftChildIndex(of index: Int) -> Int {
        return 2 * index + 1
    }
    
    private func rightChildIndex(of index: Int) -> Int {
        return 2 * index + 2
    }
}

 

위 코드에서는 Heap의 기본 동작을 구현했습니다.

 

• push(_:): 새로운 요소를 힙에 추가하고 힙의 속성을 유지합니다. 시간 복잡도는 O(log n)입니다.
• pop(): 힙의 최상위 요소를 제거하고 반환하며, 제거 후 힙의 속성을 유지합니다. 시간 복잡도는 O(log n)입니다.
• peek: 최상위 요소를 제거하지 않고 반환합니다. 시간 복잡도는 O(1)입니다.
• siftUp(from:): 새로운 요소를 삽입 후 올바른 위치로 이동시킵니다.
• siftDown(from:): 최상위 요소를 제거한 후 힙 속성을 유지하도록 재배치합니다.

구현한 Heap은 다음과 같이 사용할 수 있습니다.

// 최소 힙 테스트
var minHeap = Heap<Int> { $0 < $1 }
minHeap.push(3)
minHeap.push(1)
minHeap.push(4)
minHeap.push(1)
minHeap.push(5)

print("Min Heap Test:")
while let element = minHeap.pop() {
    print(element)
}

 

 

출력

Min Heap Test:
1
1
3
4
5

 

// 최대 힙 테스트
var maxHeap = Heap<Int> { $0 > $1 }
maxHeap.push(3)
maxHeap.push(1)
maxHeap.push(4)
maxHeap.push(1)
maxHeap.push(5)

print("\nMax Heap Test:")
while let element = maxHeap.pop() {
    print(element)
}

 

출력

Max Heap Test:
5
4
3
1
1

 

 

최소 힙은 오름차순으로 요소를 반환하고, 최대 힙은 내림차순으로 요소를 반환합니다.

이 구조는 정렬 알고리즘이나 우선순위 큐 같은 다양한 응용 프로그램에서 매우 유용하게 사용될 수 있습니다.