The seed swap drum

A seed library keeps its swap packets in a drum. Members add and remove packets, and every draw has to be fair.

The problem

Build the drum. Each packet has a distinct label, the drum holds at most one of each, and four instructions arrive:

• ("stock", label) — put that packet in. Report True, or False if it was already in the drum.
• ("withdraw", label) — take that packet out. Report True, or False if it was not in the drum.
• ("draw",) — report one packet, chosen uniformly at random from those in the drum. The packet stays in.
• ("audit", trials) — draw that many times and report how many distinct labels came up. This is how the committee checks the drum is fair.

The random source is seeded with seed, so a run repeats exactly. Every instruction except audit must take constant time.

Input. ops — a list of instruction tuples. seed — an integer.

Output. A list of the answers, in order, one per instruction.

Example.

ops  = [("stock", "borlotti"), ("stock", "borlotti"), ("draw",),
        ("withdraw", "borlotti"), ("withdraw", "borlotti"),
        ("stock", "kale"), ("draw",)]
seed = 7
  ->  [True, False, "borlotti", True, False, True, "kale"]

The second stock and the second withdraw find the drum already in the state they wanted, so they report False and change nothing.

A second example, where a packet is withdrawn from the middle:

ops  = [("stock", "rocket"), ("stock", "fennel"), ("stock", "chard"),
        ("stock", "dill"), ("withdraw", "fennel"), ("audit", 300)]
seed = 1
  ->  [True, True, True, True, True, 3]

Three hundred draws over three packets turn up three labels, never fennel.

Constraints.

• 1 <= len(ops) <= 2 * 10^5
• A label is 1 to 20 characters.
• draw and audit are only sent to a non-empty drum; 1 <= trials <= 10^4.

Hints

Approach

Brute force

Keep the packets in a list. stock appends, draw picks a random index, and withdraw scans for the label and shifts everything after it down. Each withdrawal is O(n): 2 · 10⁵ instructions on a drum of 10⁵ packets run to 10¹⁰ moves.

The insight

Keep the packets in a gapless list and a map from label to its slot: the list makes a draw one index lookup, and the map turns a withdrawal into a swap with the last packet.

A uniform draw needs the packets to fill slots 0 to n-1 with no gaps, so the hole a withdrawal leaves has to be filled. Order in the drum carries no meaning, so any packet may fill it, and the last one is free to move: removing from the end of a list costs nothing. The map is what stops the withdrawal scanning — it names the slot outright.

Algorithm

1. Keep drum, a list of labels, and slot, a map from label to its index.
2. On stock, reject a label already in slot; otherwise append it and record its index.
3. On withdraw, reject a label not in slot; otherwise take its index hole, pop the last label, and if hole is still inside the shortened list, write that label into hole and update its slot.
4. On draw, take a random index in [0, len(drum)) and report that label.
5. On audit, repeat the draw and count the distinct labels seen.

Complexity

Time O(1) for stock, withdraw and draw; an audit costs O(trials). Space O(n) — every packet appears once in the list and once in the map.

Solution

"""Seed swap drum — a dense list for drawing, a map from label to slot for removing."""

import random


def solve(ops, seed):
    drum = []       # invariant: slots 0..len(drum)-1 hold exactly the stocked packets
    slot = {}       # label -> its index in drum, kept in step with every move
    dice = random.Random(seed)
    answers = []

    def draw():
        return drum[int(dice.random() * len(drum))]

    for op in ops:
        move = op[0]
        if move == "stock":
            label = op[1]
            if label in slot:
                answers.append(False)
            else:
                slot[label] = len(drum)
                drum.append(label)
                answers.append(True)
        elif move == "withdraw":
            label = op[1]
            if label not in slot:
                answers.append(False)
            else:
                hole = slot.pop(label)
                moved = drum.pop()          # always the cheap end of the list
                if hole < len(drum):        # the hole was not the end: refill it
                    drum[hole] = moved
                    slot[moved] = hole
                answers.append(True)
        elif move == "draw":
            answers.append(draw())
        else:                               # ("audit", trials)
            seen = set()
            for _ in range(op[1]):
                seen.add(draw())
            answers.append(len(seen))

    return answers

TESTS = [
    # Stocking a packet twice changes nothing; the second call reports that.
    (
        (
            [
                ("stock", "borlotti"), ("stock", "borlotti"),
                ("draw",),
                ("withdraw", "borlotti"), ("withdraw", "borlotti"),
                ("stock", "kale"), ("draw",),
            ],
            7,
        ),
        [True, False, "borlotti", True, False, True, "kale"],
    ),
    # Withdrawing from the middle must refill the hole, or the audit sees a
    # packet that is no longer in the drum.
    (
        (
            [
                ("stock", "rocket"), ("stock", "fennel"), ("stock", "chard"), ("stock", "dill"),
                ("withdraw", "fennel"),
                ("audit", 300),
            ],
            1,
        ),
        [True, True, True, True, True, 3],
    ),
    # Every stocked packet stays reachable after a run of withdrawals.
    (
        (
            [
                ("stock", "leek"), ("stock", "orach"), ("stock", "sorrel"),
                ("withdraw", "sorrel"), ("withdraw", "leek"),
                ("stock", "cress"), ("stock", "lovage"),
                ("audit", 400),
            ],
            2,
        ),
        [True, True, True, True, True, True, True, 3],
    ),
    # Withdrawing the last-added packet takes the other branch of the refill.
    (
        (
            [
                ("stock", "sage"), ("stock", "tansy"),
                ("withdraw", "tansy"),
                ("draw",),
                ("audit", 50),
            ],
            5,
        ),
        [True, True, True, "sage", 1],
    ),
    (
        ([("withdraw", "yarrow"), ("stock", "yarrow"), ("audit", 20)], 3),
        [False, True, 1],
    ),
]

Pitfalls

• Blanking the slot instead of refilling it. The list keeps its length, a draw can land on the hole, and the audit in the second example reports 4 or crashes rather than reporting 3.
• Moving the last packet without updating its slot. The map still points at its old index, so a later withdrawal of it deletes whatever sits there now and leaves a phantom label in the drum.
• Writing the moved packet back when the hole was the last slot. After the pop, hole equals the list length, so writing there re-appends the packet you just withdrew and it can be drawn again.

Variants

• The front display — a map beside a list again, but there the list's order is the whole point.