# The Shortlist

> In every field, only a few rise to the top. Keep them.

Canonical URL: <https://datadriven.io/problems/the_shortlist>

Domain: Python · Difficulty: medium · Seniority: L5

## Problem

A scoring pipeline emits records as it runs, each a dict carrying a `category` label and a numeric `value`, and downstream only wants a per-category shortlist. For every category that appears, return its `n` highest-`value` records keyed by that category, passing the original dicts through untouched. When records in a category share a value, the one that appeared earlier in `records` comes first; a category with fewer than `n` records contributes all it has, and empty input yields an empty dict.

## Worked solution and explanation

### What this really is

This is top-N-within-a-group wearing a leaderboard costume, the Python twin of `ROW_NUMBER() OVER (PARTITION BY category ORDER BY value DESC)` filtered to rank `<= n`. Everyone gets the grouping right. What actually separates candidates is the tiebreak: when two records in a category share a value, the prompt wants the one that arrived first to win, and that is free if you lean on Python's stable sort but quietly wrong the moment you reach for a heap.

---

### Break down the requirements

#### Step 1: Bucket records by category

Partition the records into buckets keyed by their 'category' field. One pass, no sorting yet.

#### Step 2: Order each bucket by value descending

Within each bucket, order records by 'value', highest first. Do this per bucket, never globally, or you mix categories at the cutoff.

#### Step 3: Take the first N per bucket

Slice each ordered bucket to its first N records. Slicing past the end is harmless in Python, so a short bucket just returns everything it has.

#### Step 4: Keep input order for ties

Use a stable ordering so equal-valued records stay in input order. This is the requirement most people skip, and it only shows up when ties land on the N boundary.

---

### The solution

**Bucket, order descending, slice top-N per category**

```python
def top_n_per_category(n, records):
    groups = {}
    for record in records:
        groups.setdefault(record["category"], []).append(record)
    result = {}
    for category, bucket in groups.items():
        ranked = sorted(bucket, key=lambda r: r["value"], reverse=True)
        result[category] = ranked[:n]
    return result
```

*setdefault buckets in one pass; sorted() is stable so ties keep input order even with reverse=True.*

> **Time and space**
>
> Time is O(m log m) dominated by the per-bucket sorts (worst case every record in one category). Space is O(m) for the buckets. Slicing is O(n) per bucket and never errors when n exceeds the bucket size.

> **Interviewers watch for**
>
> Leaning on `sorted()` being stable to satisfy the tiebreak for free, and returning the original record dicts untouched instead of rebuilding them. Both are quiet signals that you read the contract carefully.

> **Common pitfall**
>
> Reaching for `heapq.nlargest` to get the top N. It is asymptotically tempting, but it does not guarantee input order for equal keys, so it silently breaks the tie rule the prompt demands.

---

## Common follow-up questions

- How would you express this in SQL? _(Tests `ROW_NUMBER() OVER (PARTITION BY category ORDER BY value DESC)` with a WHERE clause on the rank.)_
- What happens when N is much larger than a category's record count? _(Tests that slicing beyond the bucket size just returns all elements, no error.)_
- What changes if you needed dense ranking instead of a hard top-N cut? _(Tests assigning ranks where ties share the same rank number.)_

## Related

- [All practice problems](https://datadriven.io/problems)
- [Mock interview mode](https://datadriven.io/interview/the_shortlist)
- [Python Interview Questions](https://datadriven.io/python-interview-questions)
- [Data Engineering Interview Prep Guide](https://datadriven.io/data-engineer-interview-prep)
- [Daily Challenge](https://datadriven.io/daily)

---

Source: DataDriven (https://datadriven.io). 100% free data engineering interview prep. Live code execution against Postgres 16, Python 3.11, and Spark sandboxes. No paywall, no premium tier, no signup gate.