Swift

【数据结构与算法 - Swift实现】09 - 数据结构 堆 (

2019-05-07  本文已影响0人  Lebron_James

堆,其实是一个完整的二叉树(除了叶节点外,其他节点都有左右子节点)。堆有以下两种形式:

         10 
       /   \     
      7     9  
     / \   / \  
    6   4 5   3   
         1 
       /   \     
      2     3  
     / \   / \  
    6   8 9   7   

实现

虽然堆本质上是属于二叉树,但是我们不实用节点来实现,而是使用数组。因为堆中的元素有一定的优先级顺序,可以用数组来存储元素;并且在元素的添加移除等过程中需要频繁访问特定的位置的元素,数组正适合我们的需求。

根据以上分析,我们首先定义 Heap 如下:

struct Heap<Element: Equatable> {
    private(set) var elements: [Element] = []
    private let order: (Element, Element) -> Bool
    
    init(order: @escaping (Element, Element) -> Bool) {
        self.order = order
    }
    
    var isEmpty: Bool {
        return elements.isEmpty
    }
    
    var count: Int {
        return elements.count
    }
    
    var peek: Element? {
        return elements.first
    }
}

extension Heap: CustomStringConvertible {
    var description: String {
        return elements.description
    }
}

考虑到堆有两种形式,所以定义了一个 order 闭包,后续的相关操作可以根据它来判断堆的类型。另外还添加了一些常用属性和实现了 CustomStringConvertible 协议。

数组如何表示堆

假设我们有以下堆数据结构:

        10       第一层
-----------------------
       /   \     
      7     9    第二层
-----------------------
     / \   / \  
    6   4 5   3  第三层
-----------------------
  /
 2               第四层
-----------------------

用数组可以表示如下:

索引     0      1     2     3     4     5     6     7
        +-------------------------------------------------+
        |  10  |  7  |  9  |  6  |  4  |  5  |  3  |  2   |
        +-------------------------------------------------+
        
        |第一层 |   第二层   |         第三层         |第四层 |
        |------|-----------|-----------------------|------|      

从上图我们可以总结出,如果给定一个元素的索引 i,那么:

所以,添加获取左右子节点和父节点索引的方法,如下:

func leftChildIndex(ofParentAt index: Int) -> Int {
    return 2 * index + 1
}

func rightChildIndex(ofParentAt index: Int) -> Int {
    return 2 * index + 2
}

func parentIndex(ofChildAt index: Int) -> Int {
    return (index - 1) / 2
}

移除元素

元素的移除分两种情况:1)移除根部元素;2)移除任意位置的元素。

移除根部元素

以下面的最大堆为例,移除10:

     +------+
     |  10  |  
     +------+    
       /   \     
      7     9   
     / \   / \  
    6   4 5   3
  /
 2               

为了移除这个根部元素,我们首先将根部元素与最后一个元素的位置调换,变成:

        +-----+
        |  2  |  
        +-----+    
         /   \     
        7     9   
       / \   / \  
      6   4 5   3
    /
+------+
|  10  |  
+------+    

调换成功后,就把最后一个元素移除,变成:

           2  
         /   \     
        7     9   
       / \   / \  
      6   4 5   3

这时需要检查当前的结构是否符合堆的特性。很明显,现在的根部元素是 2,小于它的子节点,所以我们把根部元素和它的左右子节点的较大者 9 位置调换,变成:

           9  
         /   \     
        7     2   
       / \   / \  
      6   4 5   3

再继续往下比较,这时2又比它的左右子节点小,继续把 2 和它的左右子节点的较大者位置 5 调换,变成:

           9  
         /   \     
        7     5   
       / \   / \  
      6   4 2   3

2 下面没有子节点了,所以这棵树最终调整完成,符合最大堆的要求。

实现代码如下:

extension Heap {
    mutating func removePeek() -> Element? {
        guard !isEmpty else { return nil }
        elements.swapAt(0, count - 1)
        defer {
            validate(from: 0)
        }
        return elements.removeLast()
    }
    
    private mutating func validateDown(from index: Int) {
        var parentIndex = index
        while true {
            let leftIndex = leftChildIndex(ofParentAt: parentIndex)
            let rightIndex = rightChildIndex(ofParentAt: parentIndex)
            var targetParentIndex = parentIndex
            
            if leftIndex < count &&
                order(elements[leftIndex], elements[targetParentIndex]) {
                targetParentIndex = leftIndex
            }
            
            if rightIndex < count &&
                order(elements[rightIndex], elements[targetParentIndex]) {
                targetParentIndex = rightIndex
            }
            
            if targetParentIndex == parentIndex {
                return
            }
            
            elements.swapAt(parentIndex, targetParentIndex)
            parentIndex = targetParentIndex
        }
    }
}

removePeek()

validate(from index: Int),while 循环里面:

移除任意元素

移除任意元素与移除根元素的思路类似,同样是把要移除的元素跟最后一个元素调换,让后在验证是否符合堆的特性。

实现代码如下:

mutating func remove(at index: Int) -> Element? {
    guard index < elements.count else { return nil }
    
    if index == elements.count - 1 {
        return elements.removeLast()
    } else {
        elements.swapAt(index, elements.count - 1)
        defer {
            validateDown(from: index)
            validateUp(from: index)
        }
        return elements.removeLast()
    }
}

private mutating func validateUp(from index: Int) {
    var childIndex = index
    var parentIndex = self.parentIndex(ofChildAt: childIndex)
    
    while childIndex > 0 &&
        order(elements[childIndex], elements[parentIndex]) {
            elements.swapAt(childIndex, parentIndex)
            childIndex = parentIndex
            parentIndex = self.parentIndex(ofChildAt: childIndex)
    }
}

remove(at index: Int)

validateUp(from index: Int)

插入元素

有了上面的代码基础,插入元素就非常简单了,实现代码如下:

mutating func insert(_ element: Element) {
    elements.append(element)
    validateUp(from: elements.count - 1)
}

直接把插入的元素拼接到最后面,然后再向上验证堆的合法性就可以了。

搜索元素

实现代码如下:

extension Heap {
    func index(of element: Element,
               searchingFrom index: Int = 0) -> Int? {
        if index >= count {
            return nil
        }
        if order(element, elements[index]) {
            return nil
        }
        if element == elements[index] {
            return index
        }
        
        let leftIndex = leftChildIndex(ofParentAt: index)
        if let i = self.index(of: element,
                              searchingFrom: leftIndex) {
            return i
        }
        
        let rightIndex = rightChildIndex(ofParentAt: index)
        if let i = self.index(of: element,
                              searchingFrom: rightIndex) {
            return i
        }
        
        return nil
    }
}

测试

堆的代码实现就完成了,我们来测试一下:

插入
var heap = Heap<Int>(order: >)
for i in 1...7 {
    heap.insert(i)
}
print(heap)

// 结果
[7, 4, 6, 1, 3, 2, 5]

对应的树结构为:

           7  
         /   \     
        4     6   
       / \   / \  
      1   3 2   5

符合堆的特性。

移除根部元素
while !heap.isEmpty {
    print(String(describing: heap.removePeek()))
}

// 结果
Optional(7)
Optional(6)
Optional(5)
Optional(4)
Optional(3)
Optional(2)
Optional(1)

从最大的元素删除到最小的元素,符合堆的特性。

移除任意元素
let index = heap.index(of: 3) ?? 0
print(String(describing: heap.remove(at: index)))
print(heap)

// 结果
Optional(3)
[7, 5, 6, 1, 4, 2]

移除元素3,并且移除后,依然符合堆的特性。

总结

堆数据结构非常适合用来跟踪最大值或者最小值,因为获取 peek 的时间复杂度为 O(1)。在搜索性能方面,时间复杂度为 O(n), 比二叉搜索树的 O(log n) 要差很多。另外堆的插入和移除的时间复杂度都为 O(log n)

完整代码 >>

参考资料

Data Structures and Algorithms in Swift --- raywenderlich.com,如果想看原版书籍,请点击链接购买。

欢迎加入我管理的Swift开发群:536353151

下一篇文章:【数据结构与算法 - Swift实现】10 - 优先队列 (Priority Queue)

上一篇下一篇

猜你喜欢

热点阅读