Fix the smoothed target encoding for zip codes
Your task
An AI target-encoded a column of 10,000 zip codes: it replaced each zip with that zip's mean conversion rate, computed over the whole dataset before the train and test split. Every row's encoding included its own label, so the holdout AUC was great and production drifted down at once. The fix fits the encoding on training rows only and smooths rare zips toward the overall rate.
An assistant wrote encode_target_mean(train_keys, train_labels, query_keys, smoothing). It is in the editor. It has a bug. Find it and fix it.
Input and output:
train_keysandtrain_labels: each training row's category and its label, 0 or 1, at the same index. The two lists have the same length.query_keys: the categories to encode, from test or production rows.smoothing: how many rows' worth of the overall rate to blend in.- Return one encoded value per query key, in order.
Rules the tests check:
- The prior is the mean of
train_labels, or 0 if there are none. - For a query key, count the training rows with exactly that key, case included, and add up their labels to get its hits.
- If the count plus
smoothingis 0 or less, the value is the prior. - Otherwise the value is the hits plus
smoothingtimes the prior, divided by the count plussmoothing. - Answers within 0.000001 of the expected values pass.
Examples
encode_target_mean(train_keys: list[str], train_labels: list[int], query_keys: list[str], smoothing: float) -> list[float]Example 1
Input
encode_target_mean(["94107", "94107", "10001", "10001"], [1, 0, 1, 1], ["94107", "10001", "99999"], 1)Output
[0.5833333333333334, 0.9166666666666666, 0.75]Example 2
Input
encode_target_mean(["a", "a", "b"], [1, 0, 1], ["a", "b", "c"], 0)Output
[0.5, 1, 0.6666666666666666]Example 3
Input
encode_target_mean(["x", "y", "y", "y", "y"], [1, 0, 0, 0, 0], ["x"], 10)Output
[0.2727272727272727]