Fix the smoothed target encoding for zip codes

Review the AI's codeEasy10 minCapital One style

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_keys and train_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 smoothing is 0 or less, the value is the prior.
  • Otherwise the value is the hits plus smoothing times the prior, divided by the count plus smoothing.
  • 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
Inputencode_target_mean(["94107", "94107", "10001", "10001"], [1, 0, 1, 1], ["94107", "10001", "99999"], 1)
Output[0.5833333333333334, 0.9166666666666666, 0.75]
Example 2
Inputencode_target_mean(["a", "a", "b"], [1, 0, 1], ["a", "b", "c"], 0)
Output[0.5, 1, 0.6666666666666666]
Example 3
Inputencode_target_mean(["x", "y", "y", "y", "y"], [1, 0, 0, 0, 0], ["x"], 10)
Output[0.2727272727272727]