[Swift] 힙(Heap) 구현해보기

이글은 개인적인 공부를 위해 작성된 글임을 밝힙니다.

이글의 모든 내용에 대한 출처는 https://jeonyeohun.tistory.com/325 이곳입니다.

 

힙

이진트리를 사용해서 반정렬 상태로 유지하는 자료구조

힙은 트리의 루트 노드에 데이터들의 최솟값 혹은 최대값을 저장합니다.

최솟값을 루트에 둘 때는 최소힙, 최대값을 루트에 두면 최대힙이라고 합니다.

 

완전 이진트리

트리를 구성하는 모든 노드가 최대 2개의 자식노드를 가지는 자료구조

완전 이진트리는 아래와 같은 특징을 가집니다.

1. 각 높이에 있는 모든 노드가 자식노드를 가져야 다음 높이에 새 노드를 추가할 수 있습니다.
2. 가장 왼쪽부터 차례대로 노드를 채워야 합니다.

 

이 배열 기반의 이진트리를 이용해서 힙을 구현하겠습니다.

힙의 이진트리는 트리를 구성하는 각 서브트리들도 모두 힙을 만족해야합니다.

 

뼈대 만들기

import Foundation
 
struct Heap<T: Comparable> {
    private var elements: [T] = []
    private let sortFunction: (T, T) -> Bool
    
    var isEmpty: Bool {
        return self.elements.count == 1
    }
    var peek: T? {
        if self.isEmpty { return nil }
        return self.elements[1]
    }
    var count: Int {
        return self.elements.count - 1
    }
    
    init(elements: [T] = [], sortFunction: @escaping (T, T) -> Bool) {
        if !elements.isEmpty {
            self.elements = [elements.first!] + elements
        } else {
            self.elements = elements
        }
        self.sortFunction = sortFunction
        if elements.count > 1 {
            self.buildHeap()
        }
    }
    
    func leftChild(of index: Int) -> Int {
        return index * 2 
    }
    func rightChild(of index: Int) -> Int {
        return index * 2 + 1
    }
    func parent(of index: Int) -> Int {
        return (index) / 2
    }
 }

elements

실제로 데이터들을 저장할 이진트리로 표현된 배열입니다.

sortFunction

최소힙, 최대힙의 기준점으로 삼을 정렬 로직을 가진 클로저입니다.

만약 저장되는 데이터가 정수고, sortFunction이 오름차순정렬이면 최소힙, 내림차순 정렬이면 최대힙이 됩니다.

isEmpty

현재 이진트리의 요소들이 1개인지 확인합니다.

leftChild, rightChild, parent

부모와 자식노드의 인덱스를 구하는 메서드

 

삽입 - O(logN)

이렇게 최대힙이 있고, 9를 삽입해보겠습니다.

이러한 과정들을 거쳐 트리가 재조정 될 것입니다.

가장 끝에 노드를 붙였으니 노드를 올라갈 수 있는 만큼 최대한 올려주어야합니다.

그래야 트리의 각 서브트리들도 최대힙을 만족하게됩니다.

전체트리와 모든 서브트리가 최대힙을 만족하는 것을 확인할 수 있습니다.

 

삽입 구현 - swimUp

삽입의 연산은

1. 트리 제일 끝에 새 노드 붙이기
2. 새 노드를 올릴 수 있는 곳까지 올리기

이 순서로 진행됩니다.

 

2번 과정은 아래와 같이 구현이 가능합니다.

mutating func swimUp(from index: Int) {
    var index = index
    while index != 1 && self.sortFunction(self.elements[index], self.elements[self.parent(of: index)]) {
        self.elements.swapAt(index, self.parent(of: index))
        index = self.parent(of: index)
    }
}

 

루트노드에 다다르거나, 현재노드가 부모노드보다 우선순위가 낮아지기 전까지 노드와 부모노드를 계속 교환해줍니다.

 

이 알고리즘은 최악의 경우에는 지금 리프노드부터 루트노드까지 거슬러 올라가며 교환을 진행하기 때문에

시간복잡도는 트리의 높이에 대한 노드개수인 O(logN)이 됩니다.

 

1번 과정과 2번과정은 합한 전체 삽입의 구현은 아래와 같습니다.

mutating func insert(node: T) {
    if self.elements.isEmpty {
        self.elements.append(node)
        return
    }
    self.elements.append(node)
    self.swimUp(from: self.elements.endIndex - 1)
}

 

삭제 - O(logN)

삭제를 위해 루트노드와 마지막 리프노드를 교환해줍니다.

최대힙이 깨졌기 때문에 트리를 재조정합니다.

루트노드와 자식노드들 중에 우선순위가 더 높은 노드를 골라 현재 루트노드와 변경해줍니다.

이렇게 하면 모든 서브트리의 최대힙을 유지하면서 루트노드를 제거할 수 있게됩니다.

 

삭제 구현 - diveDown

삭제의 연산은

1. 루트노드와 리프노드를 교환합니다.
2. 트리를 재조정합니다.

위 두가지 순서로 진행됩니다.

 

2번과정은 아래와 같이 구현합니다.

mutating func diveDown(from index: Int) {
    var higherPriority = index
    let leftIndex = self.leftChild(of: index)
    let rightIndex = self.rightChild(of: index)

    // 왼쪽 자식 노드가 더 우선순위가 높을 때
    if leftIndex < self.elements.endIndex && self.sortFunction(self.elements[leftIndex], self.elements[higherPriority]) {
        higherPriority = leftIndex
    }
    // 오른쪽 자식 노드가 더 우선순위가 높을 때
    if rightIndex < self.elements.endIndex && self.sortFunction(self.elements[rightIndex], self.elements[higherPriority]) {
        higherPriority = rightIndex
    }
    // 교환이 필요한 경우는 교환 후 서브트리로 이동
    if higherPriority != index {
        self.elements.swapAt(higherPriority, index)
        self.diveDown(from: higherPriority)
    }
}

더 이상 교환이 필요없을 때까지 왼쪽 자식노드와 오른쪽 자식노드를 현재 루트노드와 비교하면서 우선순위가 더 높은 노드와 교환을 해주고, 교환이 되면 새로운 노드를 루트노드로 삼아서 diveDown을 재귀적으로 진행합니다.

 

최악의 경우 루트노드부터 리프노드까지 교환을 진행하는 시간복잡도 O(logN)가 요구됩니다.

 

1번 과정과 2번과정을 합쳐 삭제를 아래와 같이 구현할 수 있습니다.

mutating func remove() -> T? {
    if self.isEmpty { return nil }
    self.elements.swapAt(1, elements.endIndex - 1)
    let deleted = elements.removeLast()
    self.diveDown(from: 1)

    return deleted
}

 

배열을 힙으로 - Build Heap

만약 생성자로 받은 배열에 값이 들어 있다면 배열 전체를 다시 힙으로 재구성해주어야 하기 때문에 생성자에서 buildHeap이라는 메서드를 호출합니다.

mutating func buildHeap() {
    for index in (1...(self.elements.count / 2)).reversed() {
        self.diveDown(from: index)
    }
}

 

전체 코드

import Foundation

struct Heap<T: Comparable> {
    private var elements: [T] = []
    private let sortFunction: (T, T) -> Bool
    
    var isEmpty: Bool {
        return self.elements.count == 1
    }
    var peek: T? {
        if self.isEmpty { return nil }
        return self.elements.last!
    }
    var count: Int {
        return self.elements.count - 1
    }
    
    init(elements: [T] = [], sortFunction: @escaping (T, T) -> Bool) {
        if !elements.isEmpty {
            self.elements = [elements.first!] + elements
        } else {
            self.elements = elements
        }
        self.sortFunction = sortFunction
        if elements.count > 1 {
            self.buildHeap()
        }
    }
    
    func leftChild(of index: Int) -> Int {
        return index * 2 
    }
    func rightChild(of index: Int) -> Int {
        return index * 2 + 1
    }
    func parent(of index: Int) -> Int {
        return (index) / 2
    }
    mutating func add(element: T) {
        self.elements.append(element)
    }
    mutating func diveDown(from index: Int) {
        var higherPriority = index
        let leftIndex = self.leftChild(of: index)
        let rightIndex = self.rightChild(of: index)
        
        if leftIndex < self.elements.endIndex && self.sortFunction(self.elements[leftIndex], self.elements[higherPriority]) {
            higherPriority = leftIndex
        }
        if rightIndex < self.elements.endIndex && self.sortFunction(self.elements[rightIndex], self.elements[higherPriority]) {
            higherPriority = rightIndex
        }
        if higherPriority != index {
            self.elements.swapAt(higherPriority, index)
            self.diveDown(from: higherPriority)
        }
    }
    mutating func swimUp(from index: Int) {
        var index = index
        while index != 1 && self.sortFunction(self.elements[index], self.elements[self.parent(of: index)]) {
            self.elements.swapAt(index, self.parent(of: index))
            index = self.parent(of: index)
        }
    }
    mutating func buildHeap() {
        for index in (1...(self.elements.count / 2)).reversed() {
            self.diveDown(from: index)
        }
    }
    mutating func insert(node: T) {
        if self.elements.isEmpty {
            self.elements.append(node)
        }
        self.elements.append(node)
        self.swimUp(from: self.elements.endIndex - 1)
    }
    mutating func remove() -> T? {
        if self.isEmpty { return nil }
        self.elements.swapAt(1, elements.endIndex - 1)
        let deleted = elements.removeLast()
        self.diveDown(from: 1)
        
        return deleted
    }
}

 

---

간소화 코드

struct Heap<T: Comparable> {
    private var elements: [T] = []
    private let comparer: (T, T) -> Bool
    
    init(comparer: @escaping (T,T) -> Bool) {
        self.comparer = comparer
    }
    
    mutating func insert(element: T) {
        if elements.isEmpty {
            elements.append(element)
            elements.append(element)
            return
        }
        elements.append(element)
        swimUp(index: elements.count - 1)
    }
    
    mutating private func swimUp(index: Int) {
        var index = index
        while index > 1 && comparer(elements[index], elements[index / 2]) {
            elements.swapAt(index, index / 2)
            index /= 2
        }
    }

    mutating func pop() -> T? {
        if elements.count < 2 { return nil }
        elements.swapAt(1, elements.count - 1)
        let deletedElement = elements.removeLast()
        diveDown(index: 1)
        return deletedElement
    }
    
    mutating private func diveDown(index: Int) {
        var swapIndex = index
        var isSwap = false
        let leftIndex = index * 2
        let rightIndex = index * 2 + 1

        if leftIndex < elements.endIndex && comparer(elements[leftIndex], elements[swapIndex]) {
            swapIndex = leftIndex
            isSwap = true
        }
        
        if rightIndex < elements.endIndex && comparer(elements[rightIndex], elements[swapIndex]) {
            swapIndex = rightIndex
            isSwap = true
        }

        if isSwap {
            elements.swapAt(swapIndex, index)
            diveDown(index: swapIndex)
        }
    }
}

var maxHeap = Heap<Int>(comparer: >)
var minHeap = Heap<Int>(comparer: <)

출처 : https://dev-mandos.tistory.com/246 (만도스님 티스토리)