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. private var elements = [Element]()
  12. init(hasHigherPriority: @escaping (Element, Element) -> Bool, isEqual: @escaping (Element, Element) -> Bool) {
  13. self.hasHigherPriority = hasHigherPriority
  14. self.isEqual = isEqual
  15. }
  16. mutating func enqueue(_ element: Element) {
  17. elements.append(element)
  18. bubbleToHigherPriority(elements.count - 1)
  19. }
  20. func peek() -> Element? {
  21. elements.first
  22. }
  23. var isEmpty: Bool {
  24. 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 self.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 self.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 && self.hasHigherPriority(elements[leftChildIndex], elements[highestPriorityIndex]) {
  72. highestPriorityIndex = leftChildIndex
  73. }
  74. if rightChildIndex < elements.count && self.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. elements.debugDescription
  86. }
  87. }