首页 > 解决方案 > 如何在 Swift 中使用最小堆实现 PriorityQueue?

问题描述

首先,在 Swift 标准库中是否有内置的类似“PriorityQueue”的类?我找不到,所以我自己写。

当我从队列中弹出项目时,我希望该项目目前是最小的。对于大多数项目,这是真的,但其中一些(很少)不在正确的位置,我似乎无法找出原因。下面的代码片段可以复制粘贴到 Playground 中并运行(包括测试)。

class PriorityQueue<T> {
    private var arr: [T] = []
    private let compare: (_ a: T, _ b: T) -> Bool

    init(_ compare: @escaping (_ a: T, _ b: T) -> Bool) {
        self.compare = compare
    }

    func insert(_ a: T) {
        arr.insert(a, at: 0)
        heapify(0)
    }

    func pop() -> T? {
        if arr.isEmpty {
            return nil
        }

        let res: T = arr.removeFirst()

        if !arr.isEmpty {
            arr.insert(arr.removeLast(), at: 0)
            heapify(0)
        }

        return res
    }

    private func heapify(_ i: Int) {
        let l: Int = 2 * i + 1, r: Int = 2 * i + 2
        var smallest: Int = i

        if l < arr.count, compare(arr[l], arr[smallest]) {
            smallest = l
        }

        if r < arr.count, compare(arr[r], arr[smallest]) {
            smallest = r
        }

        if smallest == i {
            return
        }

        let temp: T = arr[i]
        arr[i] = arr[smallest]
        arr[smallest] = temp

        heapify(smallest)
    }
}

// Test case
let size: Int = 15
var arr: [Int] = Array(repeating: 0, count: size)
let queue: PriorityQueue<Int> = PriorityQueue { $0 < $1 }

for i in 0..<size {
    arr[i] = Int.random(in: 1...10000)
    queue.insert(arr[i])
}

let sorted: [Int] = arr.sorted()
for i in 0..<size {
    let n: Int = queue.pop() ?? 0
    print(n, sorted[i], n == sorted[i])
}

95% 的情况下,项目以正确的顺序从队列中弹出,但偶尔会得到如下结果:

579 579 真实

1372 1372 真

1762 1762 真

2113 2113 真

2332 2332 真

2562 2418 错误

2701 2562 错误

3137 2701 错误

2418 3137 假

4615 4615 真

6085 6085 真

6820 6820 真

7382 7382 真

8878 8878 真

9220 9220 真

标签: swiftalgorithmheappriority-queue

解决方案


将项目添加到二叉堆非常简单。您将项目附加到数组的末尾并将其在堆中冒泡到正确的位置。我不知道 Swift,所以我会给你伪代码。假设调用了支持数组(或列表)arr

insert (val)
    append val to arr
    i = arr.count - 1
    while (i > 0)
        parent = (i-1)/2
        if (arr[i] < arr[parent])
            swap(arr[i], arr[parent])
            i = parent
        else
            break

推荐阅读