1
public struct SwiftPriorityQueue<Element> {
2
    private let _hasHigherPriority: (Element, Element) -> Bool
3
    private let _isEqual: (Element, Element) -> Bool
4

5
    private var _elements = [Element]()
6

7
    public init(hasHigherPriority: @escaping (Element, Element) -> Bool, isEqual: @escaping (Element, Element) -> Bool) {
8
        _hasHigherPriority = hasHigherPriority
9
        _isEqual = isEqual
10
    }
11

12
    public mutating func enqueue(_ element: Element) {
13
        _elements.append(element)
14
        bubbleToHigherPriority(_elements.count - 1)
15
    }
16

17
    public func peek() -> Element? {
18
        return _elements.first
19
    }
20

21
    public func array() -> [Element] {
22
        return _elements
23
    }
24

25
    public var isEmpty: Bool {
26
        return _elements.count == 0
27
    }
28

29
    public var count: Int {
30
        return _elements.count
31
    }
32

33
    public mutating func dequeue() -> Element? {
34
        guard let front = peek() else {
35
            return nil
36
        }
37

38
        removeAt(0)
39

40
        return front
41
    }
42

43
    public mutating func remove(_ element: Element) {
44
        for i in 0 ..< _elements.count {
45
            if _isEqual(_elements[i], element) {
46
                removeAt(i)
47
                return
48
            }
49
        }
50
    }
51

52
    private mutating func removeAt(_ index: Int) {
53
        let removingLast = index == _elements.count - 1
54
        if !removingLast {
55
            _elements.swapAt(index, _elements.count - 1)
56
        }
57

58
        _ = _elements.popLast()
59

60
        if !removingLast {
61
            bubbleToHigherPriority(index)
62
            bubbleToLowerPriority(index)
63
        }
64
    }
65

66
    private mutating func bubbleToHigherPriority(_ initialUnbalancedIndex: Int) {
67
        precondition(initialUnbalancedIndex >= 0)
68
        precondition(initialUnbalancedIndex < _elements.count)
69

70
        var unbalancedIndex = initialUnbalancedIndex
71

72
        while unbalancedIndex > 0 {
73
            let parentIndex = (unbalancedIndex - 1) / 2
74
            guard _hasHigherPriority(_elements[unbalancedIndex], _elements[parentIndex]) else { break }
75
            _elements.swapAt(unbalancedIndex, parentIndex)
76
            unbalancedIndex = parentIndex
77
        }
78
    }
79

80
    private mutating func bubbleToLowerPriority(_ initialUnbalancedIndex: Int) {
81
        precondition(initialUnbalancedIndex >= 0)
82
        precondition(initialUnbalancedIndex < _elements.count)
83

84
        var unbalancedIndex = initialUnbalancedIndex
85
        while true {
86
            let leftChildIndex = unbalancedIndex * 2 + 1
87
            let rightChildIndex = unbalancedIndex * 2 + 2
88

89
            var highestPriorityIndex = unbalancedIndex
90

91
            if leftChildIndex < _elements.count && _hasHigherPriority(_elements[leftChildIndex], _elements[highestPriorityIndex]) {
92
                highestPriorityIndex = leftChildIndex
93
            }
94

95
            if rightChildIndex < _elements.count && _hasHigherPriority(_elements[rightChildIndex], _elements[highestPriorityIndex]) {
96
                highestPriorityIndex = rightChildIndex
97
            }
98

99
            guard highestPriorityIndex != unbalancedIndex else { break }
100
            _elements.swapAt(highestPriorityIndex, unbalancedIndex)
101

102
            unbalancedIndex = highestPriorityIndex
103
        }
104
    }
105
}
106

107
extension SwiftPriorityQueue : CustomDebugStringConvertible {
108
    public var debugDescription: String {
109
        return _elements.debugDescription
110
    }
111
}