算法练习: 优先队列和堆(Priority Queue & Heap)

[ShannonChenCHN]

优先队列

A priority queue is a queue where the most important element is always at the front.

The queue can be a max-priority queue (largest element first) or a min-priority queue (smallest element first).

参考：

• raywenderlich/swift-algorithm-club/Priority Queue

题目列表

堆

• 堆排序
• 703. 数据流中的第 K 大元素⭐️
• 215. 数组中的第K个最大元素⭐️
• 剑指 Offer 40. 最小的k个数⭐️
• 1046. 最后一块石头的重量
• 973. 最接近原点的 K 个点（难度中等）
• 347. 前 K 个高频元素（难度中等）
• 23. 合并K个升序链表（难度困难、分治算法）
• 295. 数据流的中位数（注：本题同 剑指 Offer 41. 数据流中的中位数）⭐️

总结

1. 需要掌握的基础知识有：
   • 什么是堆
     • 完全二叉树
     • 父比子大/小
   • 如何实现一个堆
     • 存储：一般是数组
     • 操作：insert、remove、shiftUp、shiftDown
   • 堆排序：建堆+排序
2. 遇到求 Top K 或者第 K 大的数这类问题时，都可以考虑是否用堆来实现
3. 重点掌握建堆的模板代码

[ShannonChenCHN]

理论知识：堆和堆排序

一、什么是堆？

堆是一种特殊的树。

只要满足以下两个条件，它就是一个堆：

• 堆是一个完全二叉树(除了最后一层，其他层的节点个数都是满的，最后一层的节点都靠左排列)；
• 堆中每一个节点的值都必须大于等于（或小于等于）其子树中每个节点的值。

大顶堆(最大堆)：每个节点的值都大于或等于子树中每个节点值的堆。
小顶堆(最小堆)：每个节点的值都小于或等于子树中每个节点值的堆。

二、如何实现一个堆？

要实现一个堆，我们需要知道：

• 如何存储一个堆
• 堆支持哪些操作

1. 如何存储一个堆

因为堆是一个完全二叉树，而完全二叉树通常又比较适合用数组来存储，所以我们可以用数组来存储堆。

如图所示，我们可以发现规律：对于数组中任意下标为 i 的节点，其左子节点下标正好是 i∗2，其右子节点的下标为 i∗2+1，其父节点的下标为 i/2​。

注：在本文的讲解以及代码中，我们都假设，堆中的数据是从数组下标为 1 的位置开始存储。那如果从 0 开始存储，实际上处理思路是没有任何变化的，唯一变化的就是计算子节点下标和父节点下标的公式。如果节点的下标是 i，那左子节点的下标就是 2∗i+1，右子节点的下标就是 2∗i+2，父节点的下标就是 (i−1​)/2。

2. 堆支持哪些操作

堆最核心的两个操作分别是：

• 往堆中插入一个元素
• 删除堆顶元素

(1) 往堆中插入一个元素

堆化（heapify）：当我们往堆中插入一个新元素后，这个树可能就不再满足应有的特性了，所以，我们就需要进行调整，让其重新满足堆的特性，这个过程就叫做堆化。

堆化有两种方式：

• 从下往上
• 从上往下

就从下往上的堆化来讲，以大顶堆为例，堆化就是顺着节点所在的路径，与父节点对比大小。如果不满足子节点小于或等于父节点的大小关系，我们就互换两个节点。一直重复这个过程，直到父子节点之间满足刚说的那种大小关系。

代码实现如下：

/// 大顶堆
class MaxHeap<T: Comparable> {
    var nodes: [T?] = [T?]()
    
    init() {
        nodes.append(nil)
    }
    
    /// 插入一个新节点
    func insert(_ data: T) {
        nodes.append(data)
        shiftUp()
    }
    
    /// 自下往上堆化
    func shiftUp() {
        var currentIndex = nodes.count - 1
        while currentIndex > 1 {
        	let parentIndex = currentIndex / 2
        	if nodes[currentIndex]! > nodes[parentIndex]! {
            	nodes.swapAt(currentIndex, parentIndex)
            	currentIndex = parentIndex
        	} else {
        		break
        	}
        }
    }
}

(2) 删除堆顶元素

我们把最后一个节点移到堆顶，然后利用同样的父子节点对比方法。对于不满足父子节点大小关系的，互换两个节点，并且重复进行这个过程，直到父子节点之间满足大小关系为止。这就是从上往下的堆化方法。

代码实现：

/// 大顶堆
class MaxHeap<T: Comparable> {
    var nodes: [T?] = [T?]()
    
    init() {
        nodes.append(nil)
    }
    
    /// 删除堆顶元素
    func remove() {
        guard nodes.count <= 1 else {
            return
        }
        
        // 将堆顶元素替换成最后一个元素
        nodes[1] = nodes.last!
        
        // 对新的堆顶元素进行堆化
        shiftDown()
    }
    
    /// 自上往下堆化
    func shiftDown() {
        var currentIndex = 1
        while true {
            
            // 从 根-左-右 中找出最大的元素作为新的堆顶元素
            var maxIndex = currentIndex
            let leftChildIndex = currentIndex * 2
            let rightChildIndex = currentIndex * 2 + 1
            
            // 堆顶元素比左子节点小
            if leftChildIndex <= nodes.count - 1 &&
                nodes[currentIndex]! < nodes[leftChildIndex]! {
                maxIndex = leftChildIndex
            }
            
            // 堆顶元素比右子节点小
            if rightChildIndex <= nodes.count - 1 &&
                nodes[maxIndex]! < nodes[rightChildIndex]! {
                maxIndex = rightChildIndex
            }
            
            // 如果当前元素本身就是最大元素，就不用继续堆化了
            if maxIndex == currentIndex {
                break
            }
            
            nodes.swapAt(maxIndex, currentIndex)
            currentIndex = maxIndex
        }
    }
}

(3) 复杂度分析

一个包含 n 个节点的完全二叉树，树的高度不会超过 log​n。堆化的过程是顺着节点所在路径比较交换的，所以堆化的时间复杂度跟树的高度成正比，也就是 O(logn)。插入数据和删除堆顶元素的主要逻辑就是堆化，所以，往堆中插入一个元素和删除堆顶元素的时间复杂度都是 O(logn)。

三、如何基于堆实现排序？

我们可以把堆排序的过程大致分解成两个大的步骤，建堆和排序。

1. 建堆

两种思路：

• 第一种思路是借助我们前面讲的，在堆中插入一个元素的思路。尽管数组中包含 n 个数据，但是我们可以假设，起初堆中只包含一个数据，就是下标为 1 的数据。然后，我们调用前面讲的插入操作，将下标从 2 到 n 的数据依次插入到堆中。这样我们就将包含 n 个数据的数组，组织成了堆。
• 第二种思路是从树的角度出发，从下往上、从后往前对每一层的节点依次进行堆化，并且每个数据都是从上往下堆化。因为叶子节点往下堆化只能自己跟自己比较，所以我们直接从最后一个非叶子节点开始，依次堆化就行了。

因为第二种方式比较常见，所以这里只讨论第二种思路。

代码实现如下：

/// 大顶堆
class MaxHeap<T: Comparable> {
    var nodes: [T?] = [T?]()
    
    init() {
        nodes.append(nil)
    }
    
    func buildHeap(_ array: [T]) {
        nodes = [nil]
        nodes += array
        
        // 从最后一个非叶子节点开始，从下向上依次堆化
        let lastParentIndex = (nodes.count - 1) / 2
        for i in stride(from: lastParentIndex, through: 1, by: -1) {
            shiftDown(i)
        }
    }
    
    /// 自上往下堆化
    func shiftDown(_ index: Int) {
        var currentIndex = index
        while true {
            
            // 从 根-左-右 中找出最大的元素作为新的堆顶元素
            var maxIndex = currentIndex
            let leftChildIndex = currentIndex * 2
            let rightChildIndex = currentIndex * 2 + 1
            
            // 堆顶元素比左子节点小
            if leftChildIndex <= nodes.count - 1 &&
                nodes[currentIndex]! < nodes[leftChildIndex]! {
                maxIndex = leftChildIndex
            }
            
            // 堆顶元素比右子节点小
            if rightChildIndex <= nodes.count - 1 &&
                nodes[maxIndex]! < nodes[rightChildIndex]! {
                maxIndex = rightChildIndex
            }
            
            // 如果当前元素本身就是最大元素，就不用继续堆化了
            if maxIndex == currentIndex {
                break
            }
            
            nodes.swapAt(maxIndex, currentIndex)
            currentIndex = maxIndex
        }
    }
}

第二种思路的时间复杂度为 O(n)。

2. 排序

建堆结束之后，数组中的数据并不是递增或者递减排序的。不过，数组中的数据已经是按照大顶堆的特性来组织的。数组中的第一个元素就是堆顶，也就是最大的元素。

我们可以把数组分成两部分，已排序部分和未排序部分。首先我们把堆顶元素跟最后一个元素交换，那最大元素就放到了下标为 n 的位置，而前面的 n - 1 个元素就是未排序部分。

接下来堆中就只剩下了未排序部分的元素，然后我们再通过堆化的方法，将剩下的 n−1 个元素重新构建成堆。堆化完成之后，我们再取堆顶的元素，放到下标为 n−1 的位置，一直重复这个过程，直到最后堆中只剩下标为 1 的一个元素，排序工作就完成了。

代码实现：

/// 大顶堆
class MaxHeap<T: Comparable> {
    var nodes: [T?] = [T?]()
    
    init() {
        nodes.append(nil)
    }
    
    /// 自上往下堆化
    func shiftDown(from index: Int, until endIndex: Int) {
        var currentIndex = index
        while true {
            
            // 从 根-左-右 中找出最大的元素作为新的堆顶元素
            var maxIndex = currentIndex
            let leftChildIndex = currentIndex * 2
            let rightChildIndex = currentIndex * 2 + 1
            
            // 堆顶元素比左子节点小
            if leftChildIndex <= endIndex &&
                nodes[currentIndex]! < nodes[leftChildIndex]! {
                maxIndex = leftChildIndex
            }
            
            // 堆顶元素比右子节点小
            if rightChildIndex <= endIndex &&
                nodes[maxIndex]! < nodes[rightChildIndex]! {
                maxIndex = rightChildIndex
            }
            
            // 如果当前元素本身就是最大元素，就不用继续堆化了
            if maxIndex == currentIndex {
                break
            }
            
            nodes.swapAt(maxIndex, currentIndex)
            currentIndex = maxIndex
        }
    }
    
    /// 排序
    func sort() {
        var lastUnsortedIndex = nodes.count - 1
        while lastUnsortedIndex > 1 {
            nodes.swapAt(lastUnsortedIndex, 1)
            
            lastUnsortedIndex -= 1
            shiftDown(from: 1, until: lastUnsortedIndex)
        }
    }
}

整个堆排序的过程，都只需要极个别临时存储空间，所以堆排序是原地排序算法。堆排序包括建堆和排序两个操作，建堆过程的时间复杂度是 O(n)，排序过程的时间复杂度是 O(nlogn)，所以，堆排序整体的时间复杂度是 O(nlogn)。

堆排序不是稳定的排序算法，因为在排序的过程，存在将堆的最后一个节点跟堆顶节点互换的操作，所以就有可能改变值相同数据的原始相对顺序。

参考

• https://time.geekbang.org/column/article/69913?utm_term=pc_interstitial_714
• https://github.com/raywenderlich/swift-algorithm-club/tree/master/Heap
• https://github.com/raywenderlich/swift-algorithm-club/tree/master/Heap%20Sort

[ShannonChenCHN]

703. 数据流中的第 K 大元素

2021.03.14 03:45 pm ~ 05:30 pm

解法一：暴力解法

每次 add 时对数组进行排序，返回第 k 个元素。

解法二：小顶堆/最小堆

思路：建小顶堆，当堆大小为 k 时，堆顶元素即为第 k 大的元素。add() 函数调用时，如果当前堆大小小于 k，就将新元素添加到最后，然后再调用 shiftUp() 进行堆化。如果当前堆大小已经达到 k 了，则判断新元素是否大于堆顶元素，大于就用新元素替换堆顶元素，然后再调用 shiftDown() 进行堆化，否则，不做任何操作。

时间复杂度：

• init 时间复杂度为：O(nlogk) ，其中 n 为初始化时 nums 的长度，logk 为每次堆化的时间复杂度。
• add 时间复杂度为：O(log k)。也就是一次堆化的时间复杂度。
  空间复杂度：O(k)。需要使用堆存储前 k 大的元素。

class KthLargest {

    private var nodes: [Int] = [Int]()
    private let capacity: Int

    init(_ k: Int, _ nums: [Int]) {
        self.capacity = k
        for num in nums {
            add(num)
        }
    }

    func add(_ val: Int) -> Int {
        if nodes.count < capacity { // size 没超过 k
            nodes.append(val)
            shiftUp()
        } else {
            if nodes.first! < val {
                nodes[0] = val
                shiftDown()
            }
        }

        return peek()
    }

    private func peek() -> Int {
        return nodes.count > 0 ? nodes.first! : 0
    }

    private func shiftUp() {
        var currentIndex = nodes.count - 1
        while currentIndex >= 1 {
            let parentIndex = (currentIndex - 1) / 2
            if nodes[currentIndex] < nodes[parentIndex] {
                nodes.swapAt(parentIndex, currentIndex)
                currentIndex = parentIndex
            } else {
                break
            }
        }
    }

    private func shiftDown() {
        var currentIndex = 0
        while currentIndex < nodes.count {
            var minIndex = currentIndex
            var leftChildIndex = currentIndex * 2 + 1
            var rightChildIndex = currentIndex * 2 + 2

            if leftChildIndex < nodes.count && nodes[leftChildIndex] < nodes[minIndex] {
                minIndex = leftChildIndex
            } 

            if rightChildIndex < nodes.count && nodes[rightChildIndex] < nodes[minIndex] {
                minIndex = rightChildIndex
            }

            if minIndex == currentIndex {
                break
            } else {
               nodes.swapAt(minIndex, currentIndex)
               currentIndex = minIndex
            }
        }
    } 
}

[ShannonChenCHN]

215. 数组中的第K个最大元素

2021.03.19 07:30 am ~ 08:30 am

解法：堆

一般有两种方式：建大顶堆依次删除 k 个元素，或者建大小为 k 的小顶堆。这里采用更常用的后者。

思路：依次将 nums 中的元素添加到堆中，在添加每个元素时，如果堆中元素个数小于 k 个，则直接插入到堆中，否则，判断新添加的元素是否比堆顶元素大，是则替换掉堆顶元素，否则 pass。

时间复杂度：O(nlogk)。堆化的时间复杂度跟树的高度成正比，也就是 O(logk)，所以，每次 insert 的时间复杂度是 O(logk)。
空间复杂度：O(k)。需要使用堆存储前 k 大的元素。

class Solution {

    func findKthLargest(_ nums: [Int], _ k: Int) -> Int {
        let heap: MinHeap = MinHeap(nums, k)
        return heap.peek()
    }
}


class MinHeap {

    private var nodes: [Int] = [Int]()
    private let capacity: Int

    init(_ nums: [Int], _ k: Int) {
        self.capacity = k
        for num in nums {
            insert(num)
        }
    }

    func peek() -> Int {
        return nodes.count > 0 ? nodes.first! : 0
    }

    func insert(_ num: Int) {

        if nodes.count < capacity { // size 没超过 k
            nodes.append(num)
            shiftUp()
        } else {
            if nodes.first! < num {
                nodes[0] = num
                shiftDown()
            }
        }
    }

    private func shiftUp() {
        var currentIndex = nodes.count - 1
        while currentIndex >= 1 {
            let parentIndex = (currentIndex - 1) / 2
            if nodes[currentIndex] < nodes[parentIndex] {
                nodes.swapAt(parentIndex, currentIndex)
                currentIndex = parentIndex
            } else {
                break
            }
        }
    }

    private func shiftDown() {
        var currentIndex = 0
        while currentIndex < nodes.count {
            var minIndex = currentIndex
            var leftChildIndex = currentIndex * 2 + 1
            var rightChildIndex = currentIndex * 2 + 2

            if leftChildIndex < nodes.count && nodes[leftChildIndex] < nodes[minIndex] {
                minIndex = leftChildIndex
            } 

            if rightChildIndex < nodes.count && nodes[rightChildIndex] < nodes[minIndex] {
                minIndex = rightChildIndex
            }

            if minIndex == currentIndex {
                break
            } else {
               nodes.swapAt(minIndex, currentIndex)
               currentIndex = minIndex
            }
        }
    } 
}

[ShannonChenCHN]

剑指 Offer 40. 最小的k个数

2021.03.20 04:30 pm ~ 05:00 pm

测试 case

输入：[2,1,3] 2
输出：[2,1]或者[1,2]

输入：[2,1,3] 0
输出：[]

解法：大顶堆

思路：构建一个大小为 k 的大顶堆，依次向堆中添加所有元素，每次添加元素时，如果堆中节点个数小于 k 个，则直接添加并堆化即可；如果堆中节点个数已经达到 k 个，则判断待添加的元素是否小于堆顶元素，是则替换堆顶元素并堆化，否则直接 pass。

时间复杂度：O(nlogk)。其中 n 是数组 arr 的长度，因为每次往堆中假如新元素时时都需要堆化，堆化的时间复杂度跟树的高度成正比，也就是 O(logk)，所以，每次 add 的时间复杂度是 O(logk)。最坏情况下数组里 n 个数都会插入，所以一共需要 O(nlog k) 的时间复杂度。
空间复杂度：O(k)。需要使用堆存储前 k 大的元素。

class Solution {
    var nodes: [Int] = [Int]()
    var capacity: Int = 0

    func getLeastNumbers(_ arr: [Int], _ k: Int) -> [Int] {
        guard k > 0 else {
            return []
        }

        capacity = k
        for num in arr {
            add(num)
        }
        return nodes
    }

    func add(_ num: Int) {
        if nodes.count < capacity {
            nodes.append(num)
            shiftUp()
        } else {
            if nodes.first! > num {
                nodes[0] = num
                shiftDown()
            }
        }
    }

    func shiftUp() {
        var currentIndex = nodes.count - 1
        while currentIndex >= 1 {
            let parentIndex = (currentIndex - 1) / 2
            if nodes[currentIndex] > nodes[parentIndex] {
                nodes.swapAt(currentIndex, parentIndex)
                currentIndex = parentIndex
            } else {
                break
            }
        }
    }

    func shiftDown() {
        var currentIndex = 0
        while currentIndex < nodes.count {
            var maxIndex = currentIndex
            let leftChildIndex = currentIndex * 2 + 1
            let rightChildIndex = currentIndex * 2 + 2

            if leftChildIndex < nodes.count && nodes[leftChildIndex] > nodes[maxIndex] {
                maxIndex = leftChildIndex
            }

            if rightChildIndex < nodes.count && nodes[rightChildIndex] > nodes[maxIndex] {
                maxIndex = rightChildIndex
            }

            if maxIndex == currentIndex {
                break
            } else {
                nodes.swapAt(maxIndex, currentIndex)
                currentIndex = maxIndex
            }
        }
    }
}

[ShannonChenCHN]

295. 数据流的中位数

注：本题同 剑指 Offer 41. 数据流中的中位数