Least-loaded task assignment with worker specialties: Frontier Interview's solution

Question: Least-loaded task assignment with worker specialties

Companies: Stripe

Difficulty: Medium

Problem statement

Tasks arrive in a fixed order and must be handed out to a team of workers so that the work stays balanced.

Each task has a unique integer id, a positive integer duration, and (in Part 2) an optional required specialty. Each worker has a unique name and (in Part 2) a set of specialties.

Process the tasks one at a time, in the given order. A worker's load is the sum of the durations already assigned to it; every worker starts at 0. Give each task to the eligible worker with the lowest load. If several eligible workers share the lowest load, pick the one whose name is lexicographically smallest (plain string comparison).

@dataclass(frozen=True)
class Task:
    id: int
    duration: int
    specialty: str | None = None

def distribute(
    workers: dict[str, set[str]], tasks: list[Task]
) -> tuple[dict[str, list[int]], list[int]]: ...

Return a map from each worker name to the ids of its tasks, in assignment order (workers with no tasks map to []), and the list of task ids that could not be assigned.

Part 1. No task has a specialty, so every worker is eligible for every task.

Part 2. A task with a specialty may only go to workers whose set contains it. A task with specialty = None may go to anyone. If no worker has the required specialty, the task is not assigned, no load changes, and its id goes in the unassigned list.

Constraints

1 <= len(workers) <= 10^5, 0 <= len(tasks) <= 10^5
1 <= duration <= 10^6
Each worker has at most 10 specialties.

Example 1 (Part 1)

workers = {"bo": {}, "al": {}, "cy": {}}
tasks   = [(10, 4), (11, 2), (12, 3), (13, 1)]   # (id, duration)

10 -> al   loads al=4 bo=0 cy=0   (all tied at 0; "al" is smallest)
11 -> bo   loads al=4 bo=2 cy=0
12 -> cy   loads al=4 bo=2 cy=3
13 -> bo   loads al=4 bo=3 cy=3

result: {"al": [10], "bo": [11, 13], "cy": [12]}, unassigned []

Example 2 (Part 2)

workers = {"alice": {"billing"}, "bob": {"billing", "fraud"}, "cara": {}}
tasks   = [(1, 5, None), (2, 3, "fraud"), (3, 4, "billing"),
           (4, 2, None), (5, 1, "risk"), (6, 6, None)]

1 -> alice  (everyone at 0)
2 -> bob    (only bob has fraud)
3 -> bob    (billing: alice=5, bob=3)
4 -> cara   (alice=5, bob=7, cara=0)
5 -> none   (nobody has risk)
6 -> cara   (alice=5, bob=7, cara=2)

result: {"alice": [1], "bob": [2, 3], "cara": [4, 6]}, unassigned [5]

Aim for better than scanning every worker for every task.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

Clarify first. Pin these down before coding:

Tasks are assigned in the given order, one at a time. This is an online greedy rule, not an optimal partition. Do not sort the tasks.
A worker's load is the sum of durations already assigned to it. Every worker starts at 0.
Ties on load go to the lexicographically smallest name (plain string order, so "Bob" < "alice").
Part 2 adds specialties. A task may name one required specialty. Only workers who have it can take the task, and the least-loaded rule applies among them. A task with no specialty can go to anyone.
What happens when no worker has the specialty? The usual answer is: leave it unassigned and report it. Ask; do not guess.

Brute force. For each task, scan every eligible worker and take the minimum of (load, name). That is O(W) per task and O(T * W) overall. It is correct and is a fine first version to write in 5 minutes. With W small it is also what you would ship. Say so, then improve it.

Part 1, optimal: a min-heap of (load, name). Python compares tuples left to right, so the heap top is the least-loaded worker with ties broken by name. Pop it, add the duration, push it back. That is O(log W) per task.

Part 2: one heap per pool, with lazy deletion. Make a pool for "any worker" plus one pool per specialty. A worker sits in the "any" pool and in the pool of each specialty it has. The difficulty: when a worker takes a task from the billing pool, its load also changes in the fraud pool and the "any" pool. Heaps cannot update a key in place.

The fix is lazy deletion. Each heap entry stores the load the worker had when the entry was pushed. After an assignment, push a fresh (new_load, name) into every pool the worker belongs to and leave the old entries where they are. When reading a heap top, check it against the worker's real load. If they differ, the entry is stale; pop it and look again.

Invariant. In every pool, each member has exactly one current entry, the one whose load equals its real load. Loads only grow (durations are positive), so a stale entry always has a smaller load than the current one. It can rise to the top, but it is recognised and dropped there. So the first current entry at the top is the true minimum of (load, name) over the pool's members.

Why not one heap and filter? Popping from a single global heap until you find someone with the specialty can scan most workers for a rare specialty, and you have to push the skipped ones back. Per-pool heaps keep each lookup to the eligible workers only.

Code

from __future__ import annotations

import heapq
from dataclasses import dataclass
from typing import Optional


@dataclass(frozen=True)
class Task:
    id: int
    duration: int
    specialty: Optional[str] = None  # None: any worker may take it


class TaskDispatcher:
    """Least-loaded assignment with one lazy min-heap per pool.

    Pool None holds every worker; pool "billing" holds the workers with
    that specialty. Heap entries are (load_when_pushed, name). An entry is
    current only while its load equals the worker's real load.
    """

    def __init__(self, workers: dict[str, set[str]]) -> None:
        if not workers:
            raise ValueError("need at least one worker")
        self.load: dict[str, int] = {name: 0 for name in workers}
        self.pools_of: dict[str, list[Optional[str]]] = {
            name: [None, *sorted(skills)] for name, skills in workers.items()
        }
        self.assigned: dict[str, list[int]] = {name: [] for name in workers}
        self.unassigned: list[int] = []
        self.heaps: dict[Optional[str], list[tuple[int, str]]] = {}
        for name, pools in self.pools_of.items():
            for pool in pools:
                self.heaps.setdefault(pool, []).append((0, name))
        for heap in self.heaps.values():
            heapq.heapify(heap)

    def _pick(self, pool: Optional[str]) -> Optional[str]:
        heap = self.heaps.get(pool, [])
        while heap:
            load, name = heap[0]
            if load == self.load[name]:
                return name
            heapq.heappop(heap)  # stale: the worker has taken work since
        return None

    def assign(self, task: Task) -> Optional[str]:
        if task.duration <= 0:
            raise ValueError(f"task {task.id}: duration must be positive")
        name = self._pick(task.specialty)
        if name is None:
            self.unassigned.append(task.id)
            return None
        self.load[name] += task.duration
        self.assigned[name].append(task.id)
        # Every old entry for this worker is now stale. Push a fresh one
        # into each pool it belongs to; the stale ones are dropped on sight.
        for pool in self.pools_of[name]:
            heapq.heappush(self.heaps[pool], (self.load[name], name))
        return name


def distribute(
    workers: dict[str, set[str]], tasks: list[Task]
) -> tuple[dict[str, list[int]], list[int]]:
    """Assign tasks in order. Returns (worker -> task ids, unassigned ids)."""
    dispatcher = TaskDispatcher(workers)
    for task in tasks:
        dispatcher.assign(task)
    return dispatcher.assigned, dispatcher.unassigned

Complexity

Let W be the number of workers, T the number of tasks, and k the largest number of specialties one worker has.

Build: O(W * k) time and space to fill and heapify the pools.
assign: O(k log P) amortized, where P is the largest heap size. Each assignment pushes k + 1 entries. Each entry is popped at most once, so the stale pops are paid for by the pushes that created them.
distribute: O(W * k + T * k * log(W + T * k)) time. Space is O(W * k + T * k) because stale entries stay until they surface. Part 1 alone (k = 0) is O((W + T) log W) with the simple pop-and-push version.
Brute force: O(T * W) time, O(W) extra space.

Test cases

Part 1. Workers bo, al, cy, tasks (10,4) (11,2) (12,3) (13,1) gives al: [10], bo: [11, 13], cy: [12]. The first three go out by name (al, bo, cy) because all loads are 0, then bo has the lowest load (2).
Part 2. Workers alice {billing}, bob {billing, fraud}, cara {}. Tasks (1,5,-) (2,3,fraud) (3,4,billing) (4,2,-) (5,1,risk) (6,6,-) gives alice: [1], bob: [2, 3], cara: [4, 6], unassigned [5].
Tie after equal loads. Workers b, a, tasks (1,2) (2,2) (3,1) gives a: [1, 3], b: [2]. After two tasks both have load 2, so a wins the tie.
One worker. Every task goes to it, in order.
No tasks. Every worker gets [], nothing is unassigned.
Specialty nobody has. The task goes to the unassigned list and no load changes.
Case matters in names: "Bob" sorts before "alice".
Duration 0 or negative raises ValueError (the lazy-deletion invariant needs loads to grow strictly).
Randomised check: compare distribute with the O(T * W) scan on a few thousand random inputs.

Follow-ups

Worker capacity. Each worker has a maximum total load. Skip a worker if the task would push it over. In the heap version, a full worker can still take a smaller task later, so you cannot just drop it. Pop candidates into a temporary list until one fits, then push them back. Ask whether a task that nobody can fit is dropped or queued.
Several required specialties per task. A task needs billing and fraud. Pools keyed by single skills no longer work. Options: key pools by the exact skill set (fine when the sets are few and known), or scan the smallest single-skill pool and filter. State the trade-off; the brute-force scan is often the right call here.
Workers join and leave, and tasks finish. Add add_worker, remove_worker, and complete(task_id) that subtracts the duration. Loads can now go down, so the "stale entries are smaller" argument fails. Keep a version counter per worker and store it in each entry; an entry is current only if its version matches. A removed worker gets a tombstone. On removal, ask whether its tasks are reassigned.
Prefer specialists. Among eligible workers, prefer the ones with fewer specialties so generalists stay free. This changes the heap key to (load, len(skills), name) or a weighted score. Make the tie-break order explicit and test it.
Optimal balance offline. If all tasks are known up front and order does not matter, minimising the maximum load is NP-hard (multiprocessor scheduling). Longest-processing-time-first (sort by duration descending, then the same greedy) is a 4/3 approximation. Mention it; do not try to solve it exactly.

Interviewer expectations

Mid-level (L4)

Gets the tie-break right and assigns tasks in the given order without sorting.
Writes the O(T * W) scan, then moves to a (load, name) heap for Part 1.
Handles Part 2 correctly, even with a filtered scan, and reports tasks nobody can take.

Senior (L5)

Asks about the no-eligible-worker case and name ordering before coding.
Spots that one assignment changes the worker's position in several pools and solves it with lazy deletion or an equivalent.
States the invariant that makes stale entries safe, and tests against the brute force.

Staff (L6+)

Explains when the brute force is the better engineering choice (small W, many specialties per task).
Knows that lazy deletion depends on loads only growing, and has a plan (versions, tombstones) once tasks complete or workers leave.
Separates the online greedy rule from the offline balancing problem and names its complexity and the standard approximation.