Swift/Swift_DataStructure

[Swift] Heap(힙) 간단 구현

Dev_Pingu 2021. 3. 16. 02:48
반응형

안녕하세요 Pingu 입니다!🐧

 

Swift에는 Heap이 따로 없어서 직접 만들어야 합니다. ㅠ.ㅠ

 

매번 만드는게 귀찮아서 이렇게 따로 글을 남기려고 합니다.

 

간단하게 배열로 힙을 만들었고 최대힙, 최소힙을 모두 만들 수 있는 힙입니다.

힙의 가장 기본적인 기능들만 구현했습니다.

 

혹시나 힙정렬이 궁금하시다면 여기를 참고해주세요.

// Made By Pingu
class Heap<T: Comparable> {
    var heapArray: [T]
    var root: T? {
        if isMaxHeap {
            maxHeapify()
        } else {
            minHeapify()
        }
        return heapArray.first
    }
    var count: Int {
        return heapArray.count
    }
    var isEmpty: Bool {
        return heapArray.isEmpty
    }
    var isMaxHeap: Bool
    
    init(_ n: [T], isMaxHeap: Bool) {
        heapArray = n
        self.isMaxHeap = isMaxHeap
    }
    func isLeftChildExist(_ parent: Int) -> Bool {
        if parent * 2 + 1 >= count {
            return false
        } else {
            return true
        }
    }
    func isRightChildExist(_ parent: Int) -> Bool {
        if parent * 2 + 2 >= count {
            return false
        } else {
            return true
        }
    }
    func parentIndex(_ child: Int) -> Int {
        return child / 2
    }
    func leftChildIndex(_ parent: Int) -> Int {
        return parent * 2 + 1
    }
    func rightChildIndex(_ parent: Int) -> Int {
        return parent * 2 + 2
    }
    func maxHeapify() {
        isMaxHeap = true
        var i = count / 2
        while i >= 0 {
            var bigChild: Int = 0
            if isLeftChildExist(i) && isRightChildExist(i) {
                if heapArray[leftChildIndex(i)] <= heapArray[rightChildIndex(i)] {
                    bigChild = rightChildIndex(i)
                } else {
                    bigChild = leftChildIndex(i)
                }
            } else if isLeftChildExist(i) && !isRightChildExist(i) {
                bigChild = leftChildIndex(i)
            } else if !isLeftChildExist(i) && !isRightChildExist(i) {
                i -= 1
                continue
            }
            
            if heapArray[bigChild] >= heapArray[i] {
                let temp = heapArray[bigChild]
                heapArray[bigChild] = heapArray[i]
                heapArray[i] = temp
            }
            i -= 1
        }
    }
    func minHeapify() {
        isMaxHeap = false
        var i = count / 2
        while i >= 0 {
            var smallChild: Int = 0
            if isLeftChildExist(i) && isRightChildExist(i) {
                if heapArray[leftChildIndex(i)] >= heapArray[rightChildIndex(i)] {
                    smallChild = rightChildIndex(i)
                } else {
                    smallChild = leftChildIndex(i)
                }
            } else if isLeftChildExist(i) && !isRightChildExist(i) {
                smallChild = leftChildIndex(i)
            } else if !isLeftChildExist(i) && !isRightChildExist(i) {
                i -= 1
                continue
            }
            
            if heapArray[smallChild] <= heapArray[i] {
                let temp = heapArray[smallChild]
                heapArray[smallChild] = heapArray[i]
                heapArray[i] = temp
            }
            i -= 1
        }
    }
    func popRoot() -> T? {
        if isMaxHeap {
            maxHeapify()
        } else {
            minHeapify()
        }
        if !heapArray.isEmpty {
            let temp = heapArray[heapArray.count - 1]
            heapArray[heapArray.count - 1] = heapArray[0]
            heapArray[0] = temp
        }
        let rootValue = heapArray.popLast()
        
        return rootValue
    }
    func push(_ n: T) {
        heapArray.append(n)
    }
}
반응형