【剑指Offer-Swift】40.最小的k个数（一题四解：暴力·快排x2·优先队列·冒泡）

Hi 👋

• Wechat: RyukieW
• 📦 技术文章归档
• 🐙 Github

我的个人项目	扫雷Elic 无尽天梯	梦见账本
类型	游戏	财务
AppStore	Elic	Umemi

欢迎关注【Lawliet的刷题小本本】专栏，和我一起刷刷刷～

虽然这题是简单难度，但是解法很多。很典型。上一题太简单了，今天就多刷一题吧～

题目

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

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

示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：

0 <= k <= arr.length <= 10000
0 <= arr[i] <= 10000

来源：力扣（LeetCode）

链接：leetcode-cn.com/problems/zu…

著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

一、 暴力解法

先排序，然后直接取。

func getLeastNumbers(_ arr: [Int], _ k: Int) -> [Int] {
    if k == 0 { return [] }
    if arr.count <= k { return arr }
    let rearr: [Int] = arr.sorted { $0 <= $1 }
    return Array(rearr[0...k-1])
}

作为优秀的程序员，虽然可以这么解，但是不应该。虽然从结果看，Swift的排序算法效率非常高。

二、 快速排序

它的基本思想是：通过一趟排序将要排序的数据分割成独立的两部分，其中一部分的所有数据都比另外一部分的所有数据都要小，然后再按此方法对这两部分数据分别进行快速排序，整个排序过程可以递归进行，以此达到整个数据变成有序序列。

2.1 辅助空间

func getLeastNumbers(_ arr: [Int], _ k: Int) -> [Int] {
    if k == 0 { return [] }
    if arr.count <= k { return arr }
    let result = quickSort(arr)
    return Array(result[0...k-1])
}

此法较为容易理解

func quickSort<T: Comparable>(_ arr: [T]) -> [T] {
    guard let base = arr.first, arr.count > 1 else {
        return arr
    }
    
    var left: [T] = []
    var right: [T] = []
    var result: [T] = []

    for index in 1..<arr.count {
        let value = arr[index]
        if value < base {
            left.append(value)
        }
        else {
            right.append(value)
        }
    }
        
    result = quickSort(left)
    result.append(base)
    result.append(contentsOf: quickSort(right))
    
    return result
}

2.2 原地快排

func getLeastNumbers(_ arr: [Int], _ k: Int) -> [Int] {
    if k == 0 { return [] }
    if arr.count <= k { return arr }
    var result = arr
    quickSort(&result, 0, arr.count - 1, k)
    return Array(result[0...k-1])
}

原地快速排序

func quickSort<T: Comparable>(_ arr: inout [T], _ low: Int, _ high: Int, _ k: Int) {
    let res = partition(&arr, low: low, high: high)
    if res == k {
        return
    }
    else if res > k {
        quickSort2(&arr, low, res-1, k)
    }
    else {
        quickSort2(&arr, res+1, high, k)
    }
}

确定基准元素的下标

func partition<T: Comparable>(_ arr: inout [T], low: Int, high: Int) -> Int {
    let root = arr[high]
    var index = low
    for i in low...high {
        if arr[i] < root {
            if i != index {
                arr.swapAt(i, index)
            }
            index = index+1
        }
    }
    
    if high != index {
        arr.swapAt(high, index)
    }
    return index
}

三、 优先队列

• 什么是优先队列
  • 按照一定条件给定的条件
  • 每次加入元素都会进行排序
  • 优先级最高的元素上浮到首位

本题我们可以：

• 限制队列元素数量
• 更大的元素具有更高优先级
• 当队列达到k值后
• 移除最高优先级的元素
• 最后剩下的就是最小的k个元素了

首先我们需要实现一下优先队列

3.1 优先队列实现

public struct SwiftPriorityQueue<Element> {
    private let _hasHigherPriority: (Element, Element) -> Bool
    private let _isEqual: (Element, Element) -> Bool

    private var _elements = [Element]()

    public init(hasHigherPriority: @escaping (Element, Element) -> Bool, isEqual: @escaping (Element, Element) -> Bool) {
        _hasHigherPriority = hasHigherPriority
        _isEqual = isEqual
    }

    public mutating func enqueue(_ element: Element) {
        _elements.append(element)
        bubbleToHigherPriority(_elements.count - 1)
    }

    public func peek() -> Element? {
        return _elements.first
    }
    
    public func array() -> [Element] {
        return _elements
    }

    public var isEmpty: Bool {
        return _elements.count == 0
    }
    
    public var count: Int {
        return _elements.count
    }

    public mutating func dequeue() -> Element? {
        guard let front = peek() else {
            return nil
        }

        removeAt(0)

        return front
    }

    public mutating func remove(_ element: Element) {
        for i in 0 ..< _elements.count {
            if _isEqual(_elements[i], element) {
                removeAt(i)
                return
            }
        }
    }

    private mutating func removeAt(_ index: Int) {
        let removingLast = index == _elements.count - 1
        if !removingLast {
            _elements.swapAt(index, _elements.count - 1)
        }

        _ = _elements.popLast()

        if !removingLast {
            bubbleToHigherPriority(index)
            bubbleToLowerPriority(index)
        }
    }

    private mutating func bubbleToHigherPriority(_ initialUnbalancedIndex: Int) {
        precondition(initialUnbalancedIndex >= 0)
        precondition(initialUnbalancedIndex < _elements.count)

        var unbalancedIndex = initialUnbalancedIndex

        while unbalancedIndex > 0 {
            let parentIndex = (unbalancedIndex - 1) / 2
            guard _hasHigherPriority(_elements[unbalancedIndex], _elements[parentIndex]) else { break }
            _elements.swapAt(unbalancedIndex, parentIndex)
            unbalancedIndex = parentIndex
        }
    }

    private mutating func bubbleToLowerPriority(_ initialUnbalancedIndex: Int) {
        precondition(initialUnbalancedIndex >= 0)
        precondition(initialUnbalancedIndex < _elements.count)

        var unbalancedIndex = initialUnbalancedIndex
        while true {
            let leftChildIndex = unbalancedIndex * 2 + 1
            let rightChildIndex = unbalancedIndex * 2 + 2

            var highestPriorityIndex = unbalancedIndex

            if leftChildIndex < _elements.count && _hasHigherPriority(_elements[leftChildIndex], _elements[highestPriorityIndex]) {
                highestPriorityIndex = leftChildIndex
            }

            if rightChildIndex < _elements.count && _hasHigherPriority(_elements[rightChildIndex], _elements[highestPriorityIndex]) {
                highestPriorityIndex = rightChildIndex
            }

            guard highestPriorityIndex != unbalancedIndex else { break }
            _elements.swapAt(highestPriorityIndex, unbalancedIndex)

            unbalancedIndex = highestPriorityIndex
        }
    }
}

extension SwiftPriorityQueue : CustomDebugStringConvertible {
    public var debugDescription: String {
        return _elements.debugDescription
    }
}

3.2 题解

func getLeastNumbers(_ arr: [Int], _ k: Int) -> [Int] {
    var stack = SwiftPriorityQueue<Int>(hasHigherPriority: { $0 > $1 }, isEqual: { $0 == $1 })
    
    arr.forEach { item in
        stack.enqueue(item)
        if stack.count > k {
            stack.dequeue()// 最高优先级的出列（最大的出去）
        }
    }
    
    return stack.array()
}

四、 逻辑冒泡

• 创建固定容量的容器
• 遍历元素
• 存满后冒泡排序
• 继续插入到头
• 遍历直到遇到大于等于自己的就停止
• 移除最后一个元素

func getLeastNumbers5(_ arr: [Int], _ k: Int) -> [Int] {
    var result: [Int] = []
    
    guard k > 0 else {
        return result
    }
    
    arr.forEach { item in
        if result.count == k, let min = result.first {
            
            if item <= min {
                result.insert(item, at: 0)
                result.removeLast()
            }
            else {
                result.insert(item, at: 0)
                
                for i in 0..<(result.count - 1) {
                    if result[i] > result[i+1] {
                        result.swapAt(i, i+1)
                    }
                    else {
                        break
                    }
                }
                
                result.removeLast()
            }
        }
        else {
            result.append(item)
            if result.count == k {
                result = result.sorted()
            }
        }
    }
    
    return result
}

总结

没想到 Swift 的 sort 函数效率还挺高的，找时间研究下具体实现。