How to implement PriorityQueue using min-heap in Swift?

Question

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

Jim Mischel · Accepted Answer · 2019-12-25 14:28:09Z

2

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

answered Dec 25, 2019 at 14:28

Jim Mischel

135k25 gold badges197 silver badges377 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

K.Wu Over a year ago

Thank you for this, Jim. (And I was wondering why my LeetCode solution can only beat 20%)

C0D3 · Accepted Answer · 2020-03-11 06:21:05Z

I'd consider writing this heapifyUp as such:

mutating private func heapifyUp(_ index: Int) {
    var childIndex = index
    let child = items[childIndex]
    var parentIndex = getParentIndex(childIndex)

    while childIndex > 0 && orderCriteria(child, items[parentIndex]) {
        items[childIndex] = items[parentIndex]
        childIndex = parentIndex
        parentIndex = getParentIndex(childIndex)
    }

    items[childIndex] = child
}

orderCriteria is simply:

private var orderCriteria: (T, T) -> Bool

And you can initialize your Heap as such:

public init(sort: @escaping (T, T) -> Bool) {
    self.orderCriteria = sort
}

That way you can use this as a MinHeap or a MaxHeap by defining the orderCriteria.

And your insert/add function becomes:

mutating public func add(_ item: T) {
    items.append(item)
    heapifyUp(items.count - 1)
}

I hope this helps! I think your solution is also very similar. Note the items I'm using is your arr and it seems your heapify() function works as well but mine is a iterative version and yours is a recursive version. Though I haven't tested yours to be honest.

K.Wu · Accepted Answer · 2019-12-22 03:48:38Z

I just figured out why I was wrong, the error occurs in func insert(_ a: T), where I simply prepend the new element and then try to heapify the array from first element.

The function should be re-written:

func insert(_ a: T) {
    arr.append(a)
    for i in stride(from: arr.count / 2 - 1, through: 0, by: -1) {
        heapify(i)
    }
}

However, this is not very efficient, as we're breaking the entire heapified array and re-heapifing what was mostly heapified. Therefore, instead of prepending the element, I append it, this way, we're breaking, at maximum, one sub-tree of the heap. From there, we try to heapify the array bottom to top:

private func heapifyUp(_ i: Int) {
    if i == 0 {
        return
    }

    let parent: Int, sibling: Int
    if i.isMultiple(of: 2) {
        parent = (i - 2) / 2
        sibling = parent * 2 + 1
    } else {
        parent = (i - 1) / 2
        sibling = parent * 2 + 2
    }

    var smallest: Int = parent

    if compare(arr[i], arr[smallest]) {
        smallest = i
    }

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

    if smallest == parent {
        return
    }

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

    heapifyUp(parent)
    heapify(smallest)
}

The mistake I made was that I assumed after I prepend an element, rest of the array (except for the new element) is still heapified, which is not true. Consider this case:

You have a heapified array: [5, 8, 6, 9, 12, 7] and when you prepend a new element 10 and call heapify(0), the array becomes: [5, 6, 8, 10, 9, 12, 7], where 8 is the parent of 12 and 7, but 8 is not smaller than 7.

You have the right idea, but your code is overly complex. See my answer.

Collectives™ on Stack Overflow

How to implement PriorityQueue using min-heap in Swift?

3 Answers 3

1 Comment

Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related