LeetCode 2532 (Hard++, Extra Category, Amazon). Time to Cross a Bridge.

Input: n = 1, k = 3, time = [[1,1,2,1],[1,1,3,1],[1,1,4,1]]
Output: 6
Explanation: 
From 0 to 1: worker 2 crosses the bridge from the left bank to the right bank.
From 1 to 2: worker 2 picks up a box from the old warehouse.
From 2 to 6: worker 2 crosses the bridge from the right bank to the left bank.
From 6 to 7: worker 2 puts a box at the new warehouse.
The whole process ends after 7 minutes. We return 6 because the problem asks for the instance of time at which the last worker reaches the left bank.

Input: n = 3, k = 2, time = [[1,9,1,8],[10,10,10,10]]
Output: 50
Explanation: 
From 0  to 10: worker 1 crosses the bridge from the left bank to the right bank.
From 10 to 20: worker 1 picks up a box from the old warehouse.
From 10 to 11: worker 0 crosses the bridge from the left bank to the right bank.
From 11 to 20: worker 0 picks up a box from the old warehouse.
From 20 to 30: worker 1 crosses the bridge from the right bank to the left bank.
From 30 to 40: worker 1 puts a box at the new warehouse.
From 30 to 31: worker 0 crosses the bridge from the right bank to the left bank.
From 31 to 39: worker 0 puts a box at the new warehouse.
From 39 to 40: worker 0 crosses the bridge from the left bank to the right bank.
From 40 to 49: worker 0 picks up a box from the old warehouse.
From 49 to 50: worker 0 crosses the bridge from the right bank to the left bank.
From 50 to 58: worker 0 puts a box at the new warehouse.
The whole process ends after 58 minutes. We return 50 because the problem asks for the instance of time at which the last worker reaches the left bank.

class Solution {

    func findCrossingTime(_ n: Int, _ k: Int, _ time: [[Int]]) -> Int {
        var lBank = PriorityQueue<Int>(comparator: { (a, b) -> Bool in
            let ta = time[a]
            let tb = time[b]
            let ca = ta[0] + ta[2]
            let cb = tb[0] + tb[2]
            if ca == cb { return b < a }  // larger index cross first
            return cb < ca  // larger cross time cross first
        })
        var rBank = PriorityQueue<Int>(comparator: { (a, b) -> Bool in
            let ta = time[a]
            let tb = time[b]
            let ca = ta[0] + ta[2]
            let cb = tb[0] + tb[2]
            if ca == cb { return b < a }  // larger index cross first
            return cb < ca  // larger cross time cross first
        })

        // 0 -> time of the worker will be waiting to cross the bridge, 1 ->idx
        var lWorker = PriorityQueue<(Int, Int)>(comparator: { (a, b) -> Bool in a.0 < b.0 })
        var rWorker = PriorityQueue<(Int, Int)>(comparator: { (a, b) -> Bool in a.0 < b.0 })

        // initially, all at left bank
        for i in 0..<k {
            lBank.add(i)
        }

        var curTime = 0
        var remainingBoxes = n
        while remainingBoxes > 0 {
            // process worker
            while !lWorker.isEmpty && lWorker.peek()!.0 <= curTime {
                lBank.add(lWorker.poll()!.1)
            }
            while !rWorker.isEmpty && rWorker.peek()!.0 <= curTime {
                rBank.add(rWorker.poll()!.1)
            }

            var worker = -1
            if !rBank.isEmpty {
                // right side can pass, a box will be put
                worker = rBank.poll()!
                let t = time[worker]
                lWorker.add((curTime + t[2] + t[3], worker))
                curTime += t[2]  // right to left

                remainingBoxes -= 1
            } else if !lBank.isEmpty && (remainingBoxes > rBank.size + rWorker.size) {
                // left side can pass
                // left side only pass when there are more boxes
                worker = lBank.poll()!
                let t = time[worker]
                rWorker.add((curTime + t[0] + t[1], worker))
                curTime += t[0]  // left to right
            } else if remainingBoxes == rBank.size + rWorker.size {
                curTime = rWorker.peek()!.0
            } else {
                // if still empty, advance time
                let nxt: Int
                if rWorker.isEmpty {
                    nxt = lWorker.peek()!.0
                } else if lWorker.isEmpty {
                    nxt = rWorker.peek()!.0
                } else {
                    nxt = min(lWorker.peek()!.0, rWorker.peek()!.0)
                }

                curTime = nxt
            }
        }

        return curTime
    }
}

// Wrapper class for PriorityQueue
struct PriorityQueue<Element> {
    private var heap: [Element]
    private let comparator: (Element, Element) -> Bool

    init(comparator: @escaping (Element, Element) -> Bool) {
        self.heap = []
        self.comparator = comparator
    }

    var isEmpty: Bool {
        return heap.isEmpty
    }

    var size: Int {
        return heap.count
    }

    func peek() -> Element? {
        return heap.first
    }

    mutating func add(_ element: Element) {
        heap.append(element)
        swim(heap.count - 1)
    }

    mutating func poll() -> Element? {
        if heap.isEmpty { return nil }
        if heap.count == 1 { return heap.removeFirst() }

        heap.swapAt(0, heap.count - 1)
        let element = heap.removeLast()
        sink(0)

        return element
    }

    private mutating func swim(_ index: Int) {
        var childIndex = index
        var parentIndex = (childIndex - 1) / 2

        while childIndex > 0 && comparator(heap[childIndex], heap[parentIndex]) {
            heap.swapAt(childIndex, parentIndex)
            childIndex = parentIndex
            parentIndex = (childIndex - 1) / 2
        }
    }

    private mutating func sink(_ index: Int) {
        var parentIndex = index

        while true {
            let leftChildIndex = 2 * parentIndex + 1
            let rightChildIndex = 2 * parentIndex + 2
            var candidateIndex = parentIndex

            if leftChildIndex < heap.count && comparator(heap[leftChildIndex], heap[candidateIndex])
            {
                candidateIndex = leftChildIndex
            }
            if rightChildIndex < heap.count
                && comparator(heap[rightChildIndex], heap[candidateIndex])
            {
                candidateIndex = rightChildIndex
            }

            if candidateIndex == parentIndex {
                return
            }

            heap.swapAt(parentIndex, candidateIndex)
            parentIndex = candidateIndex
        }
    }
}