PriorityQueue.swift 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. //
  2. // PriorityQueue.swift
  3. // Platform
  4. //
  5. // Created by Krunoslav Zaher on 12/27/15.
  6. // Copyright © 2015 Krunoslav Zaher. All rights reserved.
  7. //
  8. struct PriorityQueue<Element> {
  9. private let _hasHigherPriority: (Element, Element) -> Bool
  10. private let _isEqual: (Element, Element) -> Bool
  11. fileprivate var _elements = [Element]()
  12. init(hasHigherPriority: @escaping (Element, Element) -> Bool, isEqual: @escaping (Element, Element) -> Bool) {
  13. _hasHigherPriority = hasHigherPriority
  14. _isEqual = isEqual
  15. }
  16. mutating func enqueue(_ element: Element) {
  17. _elements.append(element)
  18. bubbleToHigherPriority(_elements.count - 1)
  19. }
  20. func peek() -> Element? {
  21. return _elements.first
  22. }
  23. var isEmpty: Bool {
  24. return _elements.count == 0
  25. }
  26. mutating func dequeue() -> Element? {
  27. guard let front = peek() else {
  28. return nil
  29. }
  30. removeAt(0)
  31. return front
  32. }
  33. mutating func remove(_ element: Element) {
  34. for i in 0 ..< _elements.count {
  35. if _isEqual(_elements[i], element) {
  36. removeAt(i)
  37. return
  38. }
  39. }
  40. }
  41. private mutating func removeAt(_ index: Int) {
  42. let removingLast = index == _elements.count - 1
  43. if !removingLast {
  44. _elements.swapAt(index, _elements.count - 1)
  45. }
  46. _ = _elements.popLast()
  47. if !removingLast {
  48. bubbleToHigherPriority(index)
  49. bubbleToLowerPriority(index)
  50. }
  51. }
  52. private mutating func bubbleToHigherPriority(_ initialUnbalancedIndex: Int) {
  53. precondition(initialUnbalancedIndex >= 0)
  54. precondition(initialUnbalancedIndex < _elements.count)
  55. var unbalancedIndex = initialUnbalancedIndex
  56. while unbalancedIndex > 0 {
  57. let parentIndex = (unbalancedIndex - 1) / 2
  58. guard _hasHigherPriority(_elements[unbalancedIndex], _elements[parentIndex]) else { break }
  59. _elements.swapAt(unbalancedIndex, parentIndex)
  60. unbalancedIndex = parentIndex
  61. }
  62. }
  63. private mutating func bubbleToLowerPriority(_ initialUnbalancedIndex: Int) {
  64. precondition(initialUnbalancedIndex >= 0)
  65. precondition(initialUnbalancedIndex < _elements.count)
  66. var unbalancedIndex = initialUnbalancedIndex
  67. while true {
  68. let leftChildIndex = unbalancedIndex * 2 + 1
  69. let rightChildIndex = unbalancedIndex * 2 + 2
  70. var highestPriorityIndex = unbalancedIndex
  71. if leftChildIndex < _elements.count && _hasHigherPriority(_elements[leftChildIndex], _elements[highestPriorityIndex]) {
  72. highestPriorityIndex = leftChildIndex
  73. }
  74. if rightChildIndex < _elements.count && _hasHigherPriority(_elements[rightChildIndex], _elements[highestPriorityIndex]) {
  75. highestPriorityIndex = rightChildIndex
  76. }
  77. guard highestPriorityIndex != unbalancedIndex else { break }
  78. _elements.swapAt(highestPriorityIndex, unbalancedIndex)
  79. unbalancedIndex = highestPriorityIndex
  80. }
  81. }
  82. }
  83. extension PriorityQueue : CustomDebugStringConvertible {
  84. var debugDescription: String {
  85. return _elements.debugDescription
  86. }
  87. }