Bypassing a gear stage

A mill drives its grindstone through a line of gear stages, and the maintenance log asks the same question about every stage at once.

The problem

The stages sit in a row along one shaft. Stage i has an integer ratio ratios[i]: the shaft leaving it turns ratios[i] times for every turn entering it. A negative ratio reverses the direction of the shaft, and a ratio of 0 means the stage is disengaged, so nothing downstream turns.

For every stage in turn, maintenance wants the ratio of the whole train with that one stage pulled out and replaced by a straight-through coupling of ratio 1 — the product of every other ratio. Answer for all stages in one go, and without division: a disengaged stage makes the whole-train product 0, and no division recovers the rest of the train from that.

Input. ratios — a list of integers, one per stage, in order along the shaft.

Output. A list bypassed of the same length, where bypassed[i] is the product of every ratio except ratios[i].

Example.

ratios = [2, 3, 4, 5]  ->  [60, 40, 30, 24]

Bypassing the first stage leaves 3 · 4 · 5 = 60; bypassing the last leaves 24.

A second example, with a disengaged stage and a reversing stage:

ratios = [4, 0, -3, 2]  ->  [0, -24, 0, 0]

Only bypassing the disengaged stage leaves the grindstone turning: 4 · (−3) · 2 = −24, negative because the reversing stage is still in the line.

Constraints.

• 2 <= len(ratios) <= 10^5
• -30 <= ratios[i] <= 30
• Every prefix product and every suffix product fits in a signed 64-bit integer.

Hints

Approach

Brute force

For each stage, multiply the other n − 1 ratios: n(n − 1) multiplications, O(n²). At 100000 stages that is about 10¹⁰ multiplications.

The insight

The train on either side of a stage never changes, so every answer is one prefix product times one suffix product, and each family is built by a single sweep.

Written out, bypassed[i] is (ratios[0] · … · ratios[i-1]) · (ratios[i+1] · … · ratios[n-1]). Multiplication is associative, so each prefix extends the one before it by a single factor, starting from the empty product, 1; the same holds backwards. And because no division appears, a disengaged stage is not a special case — it is a 0 that lands in exactly the answers containing it.

Algorithm

1. Make out of length n and set a carry to 1.
2. Walk i upward: write the carry into out[i], then fold ratios[i] into the carry.
3. Reset the carry to 1 and walk i from n − 1 down to 0, multiplying out[i] by the carry before folding ratios[i] in again.
4. Return out.

Complexity

Time O(n) — two passes, one multiplication per stage in each. Space O(1) beyond the returned list, because the output holds the prefixes while the suffixes are folded in.

Solution

"""Bypassing a gear stage — prefix and suffix products in two sweeps."""


def solve(ratios):
    n = len(ratios)
    out = [1] * n

    carry = 1
    for i in range(n):
        # invariant: carry is the product of every stage strictly before i
        out[i] = carry
        carry *= ratios[i]

    carry = 1
    for i in range(n - 1, -1, -1):
        # invariant: carry is the product of every stage strictly after i,
        # and out[i] already holds the product of everything before it
        out[i] *= carry
        carry *= ratios[i]

    return out

TESTS = [
    (([2, 3, 4, 5],), [60, 40, 30, 24]),
    (([4, 0, -3, 2],), [0, -24, 0, 0]),
    (([0, 5, 0],), [0, 0, 0]),
    (([7, 1],), [1, 7]),
    (([-2, -3, -4],), [12, 8, 6]),
    (([1, 1, 1, 1],), [1, 1, 1, 1]),
]

Pitfalls

• Dividing the whole-train product by each ratio. One disengaged stage makes that product 0: the division either raises an error or discards the only non-zero answer, losing the −24 on [4, 0, -3, 2].
• Seeding the running product with 0. The product of no stages is 1; start the carry at 0 and the whole output is zeros.
• Folding ratios[i] into the carry before writing out[i]. The carry must cover everything strictly before or after i; swap the two lines and [2, 3, 4, 5] comes back as [120, 120, 120, 120].

Variants

• Reservoir ledger — the same sweep with sums, keeping the array so later questions cost one subtraction each.
• The bracket under the shelf — needs prefix and suffix at one index too, but compares them rather than combining them.