Process Tasks Using Servers

import java.util.PriorityQueue;

class Solution {
    public int[] assignTasks(int[] servers, int[] tasks) {
        int n = servers.length, m = tasks.length;

        // Free pool: servers that can take a task now. Entry = (weight, index); top = smallest.
        PriorityQueue<int[]> freeHeap = new PriorityQueue<>((a, b) -> {
            if (a[0] != b[0]) return Integer.compare(a[0], b[0]);
            return Integer.compare(a[1], b[1]);
        });

        // Busy pool: servers still working. Entry = (freeTime, weight, index); top = earliest finish.
        PriorityQueue<long[]> busy = new PriorityQueue<>((a, b) -> {
            if (a[0] != b[0]) return Long.compare(a[0], b[0]);
            if (a[1] != b[1]) return Long.compare(a[1], b[1]);
            return Long.compare(a[2], b[2]);
        });

        for (int i = 0; i < n; i++) {
            freeHeap.offer(new int[]{servers[i], i});
        }

        int[] ans = new int[m];
        long t = 0;   // the current second

        for (int j = 0; j < m; j++) {
            // Task j is not in the queue before second j, so the clock only moves forward.
            t = Math.max(t, (long) j);

            // Move every server that has finished by second t back to the free pool.
            // Release ALL of them: several servers can finish in the same second.
            while (!busy.isEmpty() && busy.peek()[0] <= t) {
                long[] top = busy.poll();
                freeHeap.offer(new int[]{(int) top[1], (int) top[2]});
            }

            // Nobody is free: wait for the earliest finish. Jump the clock there
            // and release every server that finishes at that same second.
            if (freeHeap.isEmpty()) {
                t = busy.peek()[0];
                while (!busy.isEmpty() && busy.peek()[0] <= t) {
                    long[] top = busy.poll();
                    freeHeap.offer(new int[]{(int) top[1], (int) top[2]});
                }
            }

            int[] server = freeHeap.poll();
            int w = server[0];
            int s = server[1];
            ans[j] = s;
            // The task starts at t (not at j, which may be long before t).
            busy.offer(new long[]{t + tasks[j], (long) w, (long) s});
        }
        return ans;
    }
}

import heapq
from typing import List, Tuple

class Solution:
    def assignTasks(self, servers: List[int], tasks: List[int]) -> List[int]:
        n, m = len(servers), len(tasks)

        # Free pool: servers that can take a task now. Entry = (weight, index); top = smallest.
        free_heap: List[Tuple[int, int]] = [(servers[i], i) for i in range(n)]
        heapq.heapify(free_heap)

        # Busy pool: servers still working. Entry = (freeTime, weight, index); top = earliest finish.
        busy: List[Tuple[int, int, int]] = []

        ans = [0] * m
        t = 0   # the current second

        for j in range(m):
            # Task j is not in the queue before second j, so the clock only moves forward.
            t = max(t, j)

            # Move every server that has finished by second t back to the free pool.
            # Release ALL of them: several servers can finish in the same second.
            while busy and busy[0][0] <= t:
                ft, w, i = heapq.heappop(busy)
                heapq.heappush(free_heap, (w, i))

            # Nobody is free: wait for the earliest finish. Jump the clock there
            # and release every server that finishes at that same second.
            if not free_heap:
                t = busy[0][0]
                while busy and busy[0][0] <= t:
                    ft, w, i = heapq.heappop(busy)
                    heapq.heappush(free_heap, (w, i))

            w, s = heapq.heappop(free_heap)
            ans[j] = s
            # The task starts at t (not at j, which may be long before t).
            heapq.heappush(busy, (t + tasks[j], w, s))

        return ans

class PriorityQueue<T> {
    private heap: T[] = [];
    constructor(private compare: (a: T, b: T) => number) {}

    push(val: T): void {
        this.heap.push(val);
        this.bubbleUp(this.heap.length - 1);
    }

    pop(): T | undefined {
        if (this.heap.length === 0) return undefined;
        const top = this.heap[0];
        const bottom = this.heap.pop()!;
        if (this.heap.length > 0) {
            this.heap[0] = bottom;
            this.bubbleDown(0);
        }
        return top;
    }

    peek(): T | undefined {
        return this.heap[0];
    }

    size(): number {
        return this.heap.length;
    }

    isEmpty(): boolean {
        return this.heap.length === 0;
    }

    private bubbleUp(idx: number): void {
        while (idx > 0) {
            const p = (idx - 1) >> 1;
            if (this.compare(this.heap[idx], this.heap[p]) < 0) {
                [this.heap[idx], this.heap[p]] = [this.heap[p], this.heap[idx]];
                idx = p;
            } else {
                break;
            }
        }
    }

    private bubbleDown(idx: number): void {
        const len = this.heap.length;
        while (true) {
            let smallest = idx;
            const left = 2 * idx + 1;
            const right = 2 * idx + 2;
            if (left < len && this.compare(this.heap[left], this.heap[smallest]) < 0) {
                smallest = left;
            }
            if (right < len && this.compare(this.heap[right], this.heap[smallest]) < 0) {
                smallest = right;
            }
            if (smallest !== idx) {
                [this.heap[idx], this.heap[smallest]] = [this.heap[smallest], this.heap[idx]];
                idx = smallest;
            } else {
                break;
            }
        }
    }
}

export function assignTasks(servers: number[], tasks: number[]): number[] {
    const n = servers.length;
    const m = tasks.length;

    // Free pool: servers that can take a task now. Entry = [weight, index]; top = smallest.
    const freeHeap = new PriorityQueue<[number, number]>((a, b) => {
        if (a[0] !== b[0]) return a[0] - b[0];
        return a[1] - b[1];
    });

    // Busy pool: servers still working. Entry = [freeTime, weight, index]; top = earliest finish.
    const busy = new PriorityQueue<[number, number, number]>((a, b) => {
        if (a[0] !== b[0]) return a[0] - b[0];
        if (a[1] !== b[1]) return a[1] - b[1];
        return a[2] - b[2];
    });

    for (let i = 0; i < n; i++) {
        freeHeap.push([servers[i], i]);
    }

    const ans = new Array<number>(m);
    let t = 0;   // the current second

    for (let j = 0; j < m; j++) {
        // Task j is not in the queue before second j, so the clock only moves forward.
        t = Math.max(t, j);

        // Move every server that has finished by second t back to the free pool.
        // Release ALL of them: several servers can finish in the same second.
        while (!busy.isEmpty() && busy.peek()![0] <= t) {
            const [, w, i] = busy.pop()!;
            freeHeap.push([w, i]);
        }

        // Nobody is free: wait for the earliest finish. Jump the clock there
        // and release every server that finishes at that same second.
        if (freeHeap.isEmpty()) {
            t = busy.peek()![0];
            while (!busy.isEmpty() && busy.peek()![0] <= t) {
                const [, w, i] = busy.pop()!;
                freeHeap.push([w, i]);
            }
        }

        const [w, s] = freeHeap.pop()!;
        ans[j] = s;
        // The task starts at t (not at j, which may be long before t).
        busy.push([t + tasks[j], w, s]);
    }
    return ans;
}