Birds of a Feather

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
L4

Problem

An ingestion step hands you a list of record dicts in records and the name of a field in key. Return a dict mapping each value found under that field to the list of records carrying it, with each bucket keeping its records in the order they arrived.

Every record finds its kin; keep them in the order they arrived.

Example 1

Input:

records = [{"dept":"eng","name":"Alice"},{"dept":"sales","name":"Bob"},{"dept":"eng","name":"Carol"}], key = "dept"

Output:

{"eng":[{"dept":"eng","name":"Alice"},{"dept":"eng","name":"Carol"}],"sales":[{"dept":"sales","name":"Bob"}]}

Example 2

Input:

records = [{"id":1,"region":"west"},{"id":2,"region":"east"},{"id":3,"region":"west"}], key = "region"

Output:

{"west":[{"id":1,"region":"west"},{"id":3,"region":"west"}],"east":[{"id":2,"region":"east"}]}

Worked solution and explanation

What this really is

Strip the costume and this is GROUP BY, run in Python instead of SQL. The whole problem is 'append each record to the bucket named by its key value, making the bucket the first time you see that value.' Everyone gets the loop. What separates candidates is the two quiet decisions around it: which container avoids the if-key-not-in-dict dance, and whether the per-bucket order survives. Reach for itertools.groupby out of muscle memory and you walk straight into the trap, because groupby only groups CONSECUTIVE equal keys, so on unsorted input you get one group per run, not per value, and your buckets come out fragmented.

The build

Step 1: Pick the container that creates buckets for you

defaultdict(list) hands you an empty list the first time you touch a key, so you can append without ever checking whether the bucket exists. That is exactly the 'create bucket on first sight, extend it after' shape this needs. A plain dict with setdefault(value, []).append(record) is identical in cost and just reads noisier.

Step 2: One in-order pass

Walk records once, in order. Pull record[key], append the record to its bucket. Order is preserved for free: list.append goes to the end, and dicts keep insertion order since 3.7, so both the keys and the lists reflect the order things arrived. No sort, no second pass.

Step 3: Hand back a plain dict

Convert to dict(buckets) on the way out. A defaultdict that escapes the function is a footgun: a caller that does partitions[some_value] to check membership silently materializes an empty list and mutates your result. Returning a plain dict makes the absent key raise, which is what a caller expects.

The solution

defaultdict(list), one pass, return dict
from collections import defaultdict

def partition_by(records, key):
    buckets = defaultdict(list)
    for record in records:
        buckets[record[key]].append(record)
    return dict(buckets)
groupby (wrong here)
itertools.groupby only merges adjacent equal keys, so you must sorted(records, key=lambda r: r[key]) first. That is O(n log n) AND it destroys the original within-bucket order the spec demands.
defaultdict (correct)
A single O(n) pass appends in arrival order. Buckets are correct on unsorted input and each list stays in the order records came in.

Common follow-up questions