[Swift] (자료구조) 힙 / Heap

 

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
    }
}

 

이중 우선순위 큐(힙 구현)
MinHeap / MaxHeap

struct DualHeap<T: Comparable & Hashable> {
	struct MaxHeap<T: Comparable> {
		var heap = [T]()
		
		init(){}
		init(_ h: T) {heap.append(h); heap.append(h)}
		
		var isEmpty: Bool {heap.count <= 1}
		
		var first: T? {isEmpty ? nil : heap[1]}
		
		mutating func insert(_ h:T){
			heap.append(h)
			if isEmpty{heap.append(h)}
			var isrt = heap.count-1
			while isrt > 1 && heap[isrt] > heap[isrt/2]{
				heap.swapAt(isrt, isrt/2)
				isrt /= 2
			}
		}
		
		enum Direction {case none, lf, ryt}
		mutating func pop() -> T? {
			func haveToDown(_ pop:Int) -> Direction {
				let lf = pop*2, ryt = lf+1
				while lf >= heap.count {return .none}
				while ryt >= heap.count {return heap[lf] > heap[pop] ? .lf: .none}
				if heap[lf] <= heap[pop] && heap[ryt] <= heap[pop] {return .none}
				if heap[lf] > heap[pop] && heap[ryt] > heap[pop] {return heap[lf] > heap[ryt] ? .lf : .ryt}
				return heap[lf] > heap[pop] ? .lf: .ryt
			}
			if isEmpty {return nil}
			heap.swapAt(1, heap.count-1)
			let returnData = heap.popLast()!
			
			var pop = 1
			while true {
				switch haveToDown(pop){
				case .none:
					return returnData
				case .lf:
					heap.swapAt(pop, pop*2)
					pop *= 2
				case .ryt:
					heap.swapAt(pop, pop*2+1)
					pop = pop*2 + 1
				}
			}
		}
	}
	
	struct MinHeap<T: Comparable> {
		var heap = [T]()
		
		init(){}
		init(_ h: T) {heap.append(h); heap.append(h)}
		
		var isEmpty: Bool {heap.count <= 1}
		
		var first: T? {isEmpty ? nil : heap[1]}
		
		mutating func insert(_ h:T){
			heap.append(h)
			if isEmpty{heap.append(h)}
			var isrt = heap.count-1
			while isrt > 1 && heap[isrt] < heap[isrt/2]{
				heap.swapAt(isrt, isrt/2)
				isrt /= 2
			}
		}
		
		enum Direction { case none, lf, ryt}
		mutating func pop() -> T? {
			func haveToDown(_ pop:Int) -> Direction {
				let lf = pop*2, ryt = lf+1
				while lf >= heap.count {return .none}
				while ryt >= heap.count {return heap[lf] < heap[pop] ? .lf: .none}
				if heap[lf] >= heap[pop] && heap[ryt] >= heap[pop] {return .none}
				if heap[lf] < heap[pop] && heap[ryt] < heap[pop] {return heap[lf] < heap[ryt] ? .lf : .ryt}
				return heap[lf] < heap[pop] ? .lf: .ryt
			}
			if isEmpty {return nil}
			heap.swapAt(1, heap.count-1)
			let returnData = heap.popLast()!
			
			var pop = 1
			while true {
				switch haveToDown(pop){
				case .none:
					return returnData
				case .lf:
					heap.swapAt(pop, pop*2)
					pop *= 2
				case .ryt:
					heap.swapAt(pop, pop*2+1)
					pop = pop*2 + 1
				}
			}
		}
	}
	
	private var maxHeap = MaxHeap<T>()
	private var minHeap = MinHeap<T>()
	private var count = [T: Int]()
	
	init(){}
	
	private mutating func rmvUsed() {
		while let top = maxHeap.first {
			if count[top, default: 0] > 0 {
				break
			}
			count.removeValue(forKey: top)
			_ = maxHeap.pop()
		}
		while let top = minHeap.first {
			if count[top, default: 0] > 0 {
				break
			}
			count.removeValue(forKey: top)
			_ = minHeap.pop()
		}
	}
	
	mutating func maxFirst() -> T? {
		rmvUsed()
		return maxHeap.first
	}
	
	mutating func minFirst() -> T? {
		rmvUsed()
		return minHeap.first
	}
	
	mutating func isEmpty() -> Bool {
		rmvUsed()
		return maxHeap.isEmpty
	}
	
	mutating func insert(_ h: T) {
		maxHeap.insert(h)
		minHeap.insert(h)
		count[h, default: 0] += 1
	}
	
	mutating func maxPop() -> T? {
		rmvUsed()
		if isEmpty() { return nil }
		count[maxHeap.first!, default: 0] -= 1
		return maxHeap.pop()
	}
	
	mutating func minPop() -> T? {
		rmvUsed()
		if isEmpty() { return nil }
		count[minHeap.first!, default: 0] -= 1
		return minHeap.pop()
	}
	
	mutating func remove(_ h: T) {
		count[h, default: 0] -= 1
	}
}