PreviousNext
AI-assisted engineering/Reviewing AI code

Fix the user-level split for a recommender

Review the AI's codeEasy10 minNetflix and Spotify style

Your task

An AI split recommendation data with a plain train_test_split, row by row. Each user has many rows, so most users landed on both sides, and the model was tested on people it had already learned. Offline AUC looked great, and online results did not. The fix splits by user: all of a user's rows go to the same side.

An assistant wrote split_rows_by_user(user_ids, test_user_count). It is in the editor. It has a bug. Find it and fix it.

Input and output:

  • user_ids: the user of each row.
  • test_user_count: how many users go to the test side.
  • Return one number per row: 1 if the row is on the test side, 0 if it is in training.

Rules the tests check:

  • List the distinct users in the order they first appear.
  • The first test_user_count users in that list are test users. If test_user_count is negative, there are none. If it is more than the number of users, every user is a test user.
  • Every row of a test user is marked 1, and every other row 0.

Examples

split_rows_by_user(user_ids: list[str], test_user_count: int) -> list[int]
Example 1
Inputsplit_rows_by_user(["u1", "u2", "u1", "u3", "u2"], 1)
Output[1, 0, 1, 0, 0]
Example 2
Inputsplit_rows_by_user(["a", "b", "c"], 2)
Output[1, 1, 0]