Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import Foundation
- public struct Heap<T> {
- private var nodes = [T]()
- private var orderCriteria: (T, T) -> Bool
- public init(sort: @escaping (T, T) -> Bool) {
- self.orderCriteria = sort
- }
- public init(array: [T], sort: @escaping (T, T) -> Bool) {
- self.orderCriteria = sort
- configureHeap(from: array)
- }
- private mutating func configureHeap(from array: [T]) {
- nodes = array
- for i in (0...(nodes.count/2-1)).reversed() {
- shiftDown(i)
- }
- }
- public var isEmpty: Bool {
- return nodes.isEmpty
- }
- public var count: Int {
- return nodes.count
- }
- @inline(__always) private func parentIndex(of index: Int) -> Int {
- return (index - 1) / 2
- }
- @inline(__always) private func leftIndex(of index: Int) -> Int {
- return index * 2 + 1
- }
- @inline(__always) private func rightIndex(of index: Int) -> Int {
- return index * 2 + 2
- }
- public func peek() -> T? {
- return nodes.first
- }
- public mutating func insert(_ value: T) {
- nodes.append(value)
- shiftUp(nodes.count-1)
- }
- public mutating func insert<S: Sequence>(_ s: S) where S.Element == T {
- for value in s {
- insert(value)
- }
- }
- @discardableResult
- public mutating func remove() -> T? {
- guard !nodes.isEmpty else {
- return nil
- }
- if nodes.count == 1 {
- return nodes.removeLast()
- } else {
- // use the last node the replace the first and shift down from the first
- let value = nodes[0]
- nodes[0] = nodes.removeLast()
- shiftDown(0)
- return value
- }
- }
- fileprivate mutating func replace(index i: Int, value: T) {
- guard i < nodes.count else { return }
- remove(at: i)
- insert(value)
- }
- @discardableResult public mutating func remove(at index: Int) -> T? {
- guard index < nodes.count else {
- return nil
- }
- let size = nodes.count - 1
- if index != size {
- nodes.swapAt(index, size)
- shiftDown(from: index, until: size)
- shiftUp(index)
- }
- return nodes.removeLast()
- }
- private mutating func shiftUp(_ index: Int) {
- var childIndex = index
- let child = nodes[childIndex]
- var parentIndex = self.parentIndex(of: index)
- while childIndex > 0 && orderCriteria(child, nodes[parentIndex]) {
- nodes[childIndex] = nodes[parentIndex]
- childIndex = parentIndex
- parentIndex = self.parentIndex(of: childIndex)
- }
- nodes[childIndex] = child
- }
- private mutating func shiftDown(from index: Int, until endIndex: Int) {
- let leftChildIndex = leftIndex(of: index)
- let rightChildIndex = leftChildIndex + 1
- var first = index
- if leftChildIndex < endIndex && orderCriteria(nodes[leftChildIndex], nodes[first]) {
- first = leftChildIndex
- }
- if rightChildIndex < endIndex && orderCriteria(nodes[rightChildIndex], nodes[first]) {
- first = rightChildIndex
- }
- if first == index {
- return
- }
- nodes.swapAt(index, first)
- shiftDown(from: first, until: endIndex)
- }
- private mutating func shiftDown(_ index: Int) {
- shiftDown(from: index, until: nodes.count)
- }
- }
- extension Heap {
- public mutating func sort() -> [T] {
- for i in stride(from: (nodes.count - 1), through: 1, by: -1) {
- nodes.swapAt(0, i)
- shiftDown(from: 0, until: i)
- }
- return nodes
- }
- }
- public func heapsort<T>(_ a: [T], _ sort: @escaping (T, T) -> Bool) -> [T] {
- let reverseOrder = { i1, i2 in sort(i2, i1) }
- var h = Heap(array: a, sort: reverseOrder)
- return h.sort()
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement