Even dye lots

A dye works records the dye lot used on each bolt of cloth. Find the longest opening run that pulling one bolt makes even.

The problem

The day's sheet is a list of dye codes, one per bolt, in the order the bolts went through the vat. A run is even when every dye still present in it has been used on the same number of bolts as every other.

The foreman may pull exactly one bolt — not zero, not two — and send it back for re-dyeing. A dye whose only bolt is pulled stops appearing, and an absent dye places no demand on the count. Find the longest prefix of bolts that one pull makes even.

Input. bolts — a list of integers, the dye code used on each bolt, in order.

Output. An integer: the length of that prefix.

Example.

bolts = [5, 5, 3, 3, 8, 1, 1, 8]   ->  7

The first seven bolts carry dye 5 twice, 3 twice, 1 twice and 8 once. Pull the lone 8 and three dyes are left on two bolts each. All eight fail: four dyes sit on two bolts, and any pull drops one of them to one.

A second example, where the whole run is one dye:

bolts = [6, 6, 6, 6]           ->  4
bolts = [7, 7, 7, 2, 2, 2, 9]  ->  7

One count cannot disagree with itself, so any bolt of dye 6 will do. In the second sheet the lone 9 is the pull.

Constraints.

• 1 <= len(bolts) <= 10^5
• 1 <= bolts[i] <= 10^5
• The pull is compulsory and must fall inside the run.

Hints

Approach

Brute force

For each of the n runs, count the uses, then try pulling a bolt of each dye present and check whether the survivors agree. Counting alone is O(n) per run: at least n² steps, around 10¹⁰ for a full day's sheet.

The insight

Keep a tally of the tallies: for each use count, how many dyes sit at it. That turns the whole test into three O(1) checks on the largest count.

Write top for the largest use count in the run. Three shapes are fixable: every dye used once, so any pull works; one dye leading by a single bolt with the rest tied one behind, so pulling one of its bolts joins the tie; one dye on a single bolt with the rest tied at top, so pulling that bolt removes it. A bolt moves one dye from one count to the next, so both maps change by a constant amount.

Algorithm

1. Keep uses[dye] and tally[count], plus top, the largest live count.
2. Per bolt, take the dye's old count c, decrement tally[c] when c > 0, set the count to c + 1 and increment tally[c + 1].
3. Raise top to c + 1 if that is larger.
4. Record the length when top == 1, or tally[top] == 1 and top + tally[top - 1] * (top - 1) equals it, or tally[1] == 1 and 1 + top * tally[top] equals it.
5. Return the last length that qualified.

Complexity

Time O(n) — one pass, constant work per bolt. Space O(k) for the two maps, k being the number of distinct dye codes.

Solution

"""Even dye lots — a tally of the use counts alongside the counts themselves."""

from collections import defaultdict


def solve(bolts):
    uses = defaultdict(int)       # dye code -> bolts dyed with it so far
    tally = defaultdict(int)      # use count -> how many dyes sit at it
    top = 0                       # invariant: top is the largest live use count
    best = 0

    for length, dye in enumerate(bolts, 1):
        before = uses[dye]
        if before:
            tally[before] -= 1
        uses[dye] = before + 1
        tally[before + 1] += 1
        if before + 1 > top:
            top = before + 1

        # A run is fixable in exactly three shapes:
        #   every dye used once             -> pull any bolt
        #   one dye leads by a single bolt  -> pull one of its bolts
        #   exactly one dye used once       -> pull that bolt entirely
        all_singles = top == 1
        lone_leader = tally[top] == 1 and top + tally[top - 1] * (top - 1) == length
        lone_stray = tally[1] == 1 and 1 + top * tally[top] == length
        if all_singles or lone_leader or lone_stray:
            best = length

    return best

TESTS = [
    (([5, 5, 3, 3, 8, 1, 1, 8],), 7),
    (([7, 7, 7, 2, 2, 2, 9],), 7),
    # One dye all the way: pulling a bolt leaves a single count, which is even.
    (([6, 6, 6, 6],), 4),
    (([4, 1, 9, 3],), 4),
    (([9],), 1),
    (([2, 2, 4, 4, 6, 6],), 5),
    (([1, 1, 1, 2, 2, 2],), 5),
]

Pitfalls

• Rejecting a run that is one dye repeated. For [6, 6, 6, 6] the leader test passes with tally[top - 1] at zero — 4 + 0 == 4. Insisting on other dyes at top - 1 returns 1 instead of 4.
• Treating an already-even run as an answer. All eight bolts of the first example are balanced, and that is why they fail: the pull is compulsory, and it unbalances them.
• Forgetting to decrement the old count. The stale entry leaves tally[top] too large, the leader test never fires, and long runs of one dye report a short answer.

Variants

• The seed swap drum — a second structure kept exactly in step with the first, there for O(1) removal rather than O(1) counting.