The Shortlist

A medium Python interview practice problem on DataDriven. Write the Python and run it against test cases, with instant feedback.

Domain
Python
Difficulty
medium
Seniority
L5
Asked in
Meta

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.

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

Example 1

Input:

n = 2, records = [{"value":10,"category":"A"},{"value":30,"category":"A"},{"value":20,"category":"B"},{"value":20,"category":"A"},{"value":50,"category":"B"},{"value":40,"category":"B"}]

Output:

{"A":[{"value":30,"category":"A"},{"value":20,"category":"A"}],"B":[{"value":50,"category":"B"},{"value":40,"category":"B"}]}

Example 2

Input:

n = 3, records = []

Output:

{}

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
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.

Common follow-up questions