Fix the cohort retention check for a dashboard query
Your task
An agent wrote a 60-line SQL query for monthly cohort retention, and nobody on the team fully follows its window functions. Before executives see the numbers, you check the query against an oracle: a second calculation, small enough to be plainly right, that the query must agree with.
An assistant wrote cohort_retention_counts(signup_lines, activity_lines, cohort_month, max_offset). It is in the editor. It has a bug. Find it and fix it.
Input and output:
signup_lines: one string per signup,"user_id|date", with dates like"2026-01-05".activity_lines: one string per activity event, in the same format.cohort_month: the signup month to study, as"YYYY-MM".max_offset: how many months after the cohort month to count, an integer.- Return a list of integers.
Rules the tests check:
- Split each signup or activity line on
|. Skip a line that does not have exactly 2 parts, or whose date part is shorter than 7 characters. - A line's month is the first 7 characters of its date. The rest of the date is not checked.
- User ids compare exactly. An empty id is still an id.
cohort_monthmust be 7 characters: 4 digits,-, then 2 digits from 01 to 12. Otherwise return an empty list. A negativemax_offsetalso returns an empty list.- The cohort is every distinct user with at least one signup line in
cohort_month. A user with several signup lines in that month is one member. - Index 0 of the result is the size of the cohort.
- For each i from 1 to
max_offset, index i is the number of distinct cohort members with at least one activity in the month i months aftercohort_month. Months roll into the next year: 1 month after"2025-12"is"2026-01". - Several activities by one member in one month count once. Activity by users outside the cohort is ignored, and so is activity in any month that is not one of the counted months.
- The result always has
max_offset + 1entries, with 0 for a month nobody was active.
Examples
cohort_retention_counts(signup_lines: list[str], activity_lines: list[str], cohort_month: str, max_offset: int) -> list[int]Example 1
Input
cohort_retention_counts(["u1|2026-01-05", "u2|2026-01-20", "u3|2026-02-02", "u4|2026-01-31"], ["u1|2026-02-11", "u1|2026-02-28", "u2|2026-03-01", "u3|2026-02-03"], "2026-01", 2)Output
[3, 1, 1]Example 2
Input
cohort_retention_counts(["a|2025-12-31", "b|2025-12-01"], ["a|2026-01-15", "a|2026-01-16", "b|2026-02-02", "a|2026-03-05", "b|2026-03-05"], "2025-12", 3)Output
[2, 1, 1, 2]