Search code examples
swiftalgorithmheappriority-queue

How to implement PriorityQueue using min-heap in Swift?


First of all, is there a built-in "PriorityQueue"-like class in Swift standard library? I am unable to find one, therefore I'm writing it myself.

When I pop items from my queue, I expect the item to be the smallest at the moment. For most items, this is true, but some of them (very few) are not in the right positions and I can't seem to find out why. The below code snippet can be copy-pasted in Playground and run (including tests).

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% of the times, items are being popped from the queue in correct order, but occasionally, I get results like this:

579 579 true

1372 1372 true

1762 1762 true

2113 2113 true

2332 2332 true

2562 2418 false

2701 2562 false

3137 2701 false

2418 3137 false

4615 4615 true

6085 6085 true

6820 6820 true

7382 7382 true

8878 8878 true

9220 9220 true


Solution

  • Adding an item to a binary heap is very simple. You append the item to the end of the array and bubble it up the heap to its proper position. I don't know Swift, so I'll give you the pseudo-code. Assuming that backing array (or list) is called 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