diff --git a/PriorityQueue/PriorityQueue.swift b/PriorityQueue/PriorityQueue.swift index df5be22..708738d 100644 --- a/PriorityQueue/PriorityQueue.swift +++ b/PriorityQueue/PriorityQueue.swift @@ -9,20 +9,20 @@ import Foundation public class PriorityQueue { - + private final var _heap: [T] private let compare: (T, T) -> Bool - - public init(_ compare: (T, T) -> Bool) { + + public init(_ compare: @escaping (T, T) -> Bool) { _heap = [] self.compare = compare } - + public func push(newElement: T) { _heap.append(newElement) - siftUp(_heap.endIndex - 1) + siftUp(index: _heap.endIndex - 1) } - + public func pop() -> T? { if _heap.count == 0 { return nil @@ -30,17 +30,17 @@ public class PriorityQueue { if _heap.count == 1 { return _heap.removeLast() } - swap(&_heap[0], &_heap[_heap.endIndex - 1]) + _heap.swapAt(0, _heap.endIndex - 1) let pop = _heap.removeLast() - siftDown(0) + siftDown(index: 0) return pop } - + private func siftDown(index: Int) -> Bool { let left = index * 2 + 1 let right = index * 2 + 2 var smallest = index - + if left < _heap.count && compare(_heap[left], _heap[smallest]) { smallest = left } @@ -48,21 +48,21 @@ public class PriorityQueue { smallest = right } if smallest != index { - swap(&_heap[index], &_heap[smallest]) - siftDown(smallest) + _heap.swapAt(index, smallest) + siftDown(index: smallest) return true } return false } - + private func siftUp(index: Int) -> Bool { if index == 0 { return false } let parent = (index - 1) >> 1 if compare(_heap[index], _heap[parent]) { - swap(&_heap[index], &_heap[parent]) - siftUp(parent) + _heap.swapAt(index, parent) + siftUp(index: parent) return true } return false @@ -73,54 +73,54 @@ extension PriorityQueue { public var count: Int { return _heap.count } - + public var isEmpty: Bool { return _heap.isEmpty } - - public func update(element: T2) -> T? { + + public func update(element: T2) -> T? where T2: Equatable { assert(element is T) // How to enforce this with type constraints? - for (index, item) in _heap.enumerate() { + for (index, item) in _heap.enumerated() { if (item as! T2) == element { _heap[index] = element as! T - if siftDown(index) || siftUp(index) { + if siftDown(index: index) || siftUp(index: index) { return item } } } return nil } - - public func remove(element: T2) -> T? { + + public func remove(element: T2) -> T? where T2: Equatable { assert(element is T) // How to enforce this with type constraints? - for (index, item) in _heap.enumerate() { + for (index, item) in _heap.enumerated() { if (item as! T2) == element { - swap(&_heap[index], &_heap[_heap.endIndex - 1]) + _heap.swapAt(index, _heap.endIndex - 1) _heap.removeLast() - siftDown(index) + siftDown(index: index) return item } } return nil } - + public var heap: [T] { return _heap } - + public func removeAll() { _heap.removeAll() } } -extension PriorityQueue: GeneratorType { +extension PriorityQueue: IteratorProtocol { public typealias Element = T public func next() -> Element? { return pop() } } -extension PriorityQueue: SequenceType { +extension PriorityQueue: Sequence { public typealias Generator = PriorityQueue public func generate() -> Generator { return self