Wristband tally

A festival hands out coloured wristbands and never says how many it printed. A marshal asks one question of a few attendees, and the answers alone pin down a floor on the crowd.

The problem

Every attendee wears one wristband, in one of several colours. The marshal stops some attendees — not all — and asks each: "how many other people here wear your colour?" Everyone asked answers truthfully, and nobody is asked twice.

From the answers, work out the smallest crowd consistent with them; attendees never asked still count if the answers require them to exist.

Two facts do the work. Attendees of the same colour give the same answer, so different answers mean different colours. And an attendee answering a sits in a colour group of exactly a + 1 — themselves plus the a others.

Input. answers — a list of non-negative integers, one per attendee asked, possibly empty.

Output. An integer: the smallest possible number of attendees.

Example.

answers = [1, 1, 2]   ->  5

The two who said 1 can share a colour, a complete group of 2. The one who said 2 needs a group of 3, so two unasked people exist. Total 5.

A second example, where one answer forces more than one group:

answers = [2, 2, 2, 2]   ->  6

A group whose members answer 2 holds exactly 3 people, so four such attendees cannot share a colour. Two groups are needed, each paid for in full: 3 + 3 = 6.

Constraints.

• 0 <= len(answers) <= 10^3
• 0 <= answers[i] <= 10^3

Hints

Approach

Brute force

Assign colours one at a time, opening a new colour when the current one fills, and search the orderings for the assignment with the fewest people. The search branches on every attendee, so it is exponential in the number asked.

The insight

Answers partition the crowd: different answers mean different colours, so each distinct answer is costed on its own and the costs added.

The answer belongs to the colour group, not the person: everyone in a group of size s answers s - 1. Within the bucket for answer a every group holds a + 1, so packing k attendees into as few groups as possible is ceiling division, and each open group is filled to a + 1 because its members insist that many share their colour. Buckets never interact, so no search is needed.

Algorithm

1. Tally the answers into a map from answer to how many attendees gave it.
2. Start the crowd at 0.
3. For each pair (a, k) in the tally, let size = a + 1.
4. Compute groups = ceil(k / size) as (k + size - 1) // size.
5. Add groups * size to the crowd; return it when the tally is done.

Complexity

Time O(n) — one pass to tally, one over at most n distinct answers. Space O(n) — one tally entry per distinct answer.

Solution

"""Wristband tally — bucket the answers, then fill each colour group to capacity."""

from collections import Counter


def solve(answers):
    crowd = 0
    for answer, asked in Counter(answers).items():
        # invariant: everyone answering `answer` sits in a group of exactly
        # answer + 1 people, and groups for different answers never mix.
        size = answer + 1
        groups = (asked + size - 1) // size      # ceiling: a part-filled group still costs `size`
        crowd += groups * size
    return crowd

TESTS = [
    (([1, 1, 2],), 5),
    (([2, 2, 2, 2],), 6),
    (([0, 0, 0],), 3),
    (([],), 0),
    (([5],), 6),
    (([1, 1, 1, 1],), 4),
    (([0, 1, 1, 2, 2, 2],), 6),
]

Pitfalls

• Using a as the group size instead of a + 1. The answer excludes the speaker, so [1, 1, 2] gives 4 rather than 5, and [0, 0, 0] gives 0 for three attendees standing right there.
• Putting every attendee with the same answer into one group. On [2, 2, 2, 2] that reports 3, and three people cannot contain four.
• Adding k rather than groups * size. The unasked members of a part-filled group vanish: [2, 2, 2, 2] reports 4, not 6.
• Rounding down. 4 // 3 is 1 group, one short. Round up with (k + size - 1) // size.

Variants

• Interfering channels — a count map again, looked up by key rather than swept.
• What a hash map buys you — the lesson on counting with a map.