Permalink
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
156 lines (127 sloc) 4.13 KB
/**
* 堆
*/
public struct Heap<T> {
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 stride(from: (nodes.count/2 - 1), through: 0, by: -1) {
shiftDown(i)
}
}
// private mutating func buildHeap(fromArray array: [T]) {
// elements = array
// for i in stride(from: (nodes.count/2-1), through: 0, by: -1) {
// shiftDown(index: i, heapSize: elements.count)
// }
// }
public var isEmpty: Bool {
return nodes.isEmpty
}
public var count: Int {
return nodes.count
}
///
@inline(__always) internal func parentIndex(ofIndex i: Int) -> Int {
return (i - 1) / 2
}
@inline(__always) internal func leftChildIndex(ofIndex i: Int) -> Int {
return 2*i + 1
}
@inline(__always) internal func rightChildIndex(ofIndex i: Int) -> Int {
return 2*i + 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>(_ sequence: S) where S.Iterator.Element == T {
for value in sequence {
insert(value)
}
}
public mutating func replace(index i: Int, value: T) {
guard i < nodes.count else {
return
}
}
/// 删除根节点
@discardableResult public mutating func remove() -> T? {
guard !nodes.isEmpty else {
return nil
}
if nodes.count == 1 {
return nodes.removeLast()
} else {
let value = nodes[0]
nodes[0] = nodes.removeLast()
shiftDown(0)
return 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()
}
/// 向上移动
internal mutating func shiftUp(_ index: Int) {
var childIndex = index
let child = nodes[childIndex]
var parentIndex = self.parentIndex(ofIndex: childIndex)
while childIndex > 0 && orderCriteria(child, nodes[parentIndex]) {
nodes[childIndex] = nodes[parentIndex]
childIndex = parentIndex
parentIndex = self.parentIndex(ofIndex: childIndex)
}
nodes[childIndex] = child
}
/// 向下移动
internal mutating func shiftDown(from index: Int, until endIndex: Int) {
let leftChildIndex = self.leftChildIndex(ofIndex: 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)
}
internal mutating func shiftDown(_ index: Int) {
shiftDown(from: index, until: nodes.count)
}
}
// MARK: - Searching
extension Heap where T: Equatable {
public func index(of node: T) -> Int? {
return nodes.index(where: { $0 == node })
}
@discardableResult public mutating func remove(node: T) -> T? {
if let index = index(of: node) {
return remove(at: index)
}
return nil
}
}