0

Factory network cost with stage-to-stage transport

by · 1 day ago· ⑂ 0 remixes

Question: Factory network cost with stage-to-stage transport

Companies: Stripe

Difficulty: Hard

Problem statement

A company is planning a chain of factories. Each factory has several candidate options, and you must pick exactly one option for every factory. The interview comes in parts; each part builds on the last.

Part 1: independent build costs

Implement min_build_cost(costs: list[list[int]]) -> int. costs[i] lists the build cost of each option for factory i. Rows may have different lengths. Return the minimum total build cost (the sum of the chosen options).

costs = [[4, 2, 7], [5], [3, 9]]
min_build_cost(costs)  # 10  (2 + 5 + 3)

Part 2: stages with transport cost

Now the factories are stages of one production line, in order. Each option is a site (position, build_cost) on a straight road. Goods built at stage i are shipped to the site chosen for stage i + 1, which costs |x - y| for positions x and y.

Implement min_network_cost(stages: list[list[tuple[int, int]]]) -> int. Return the minimum of (sum of build costs of the chosen sites) + (sum of transport costs between consecutive stages). Start with three stages.

stages = [
    [(0, 5), (10, 1)],   # stage 0
    [(3, 2), (9, 6)],    # stage 1
    [(4, 1)],            # stage 2
]
min_network_cost(stages)  # 12
# 0 -> 3 -> 4 costs 5 + 2 + 1 to build plus 3 + 1 to ship = 12.
# 10 -> 3 -> 4 also costs 12. The cheapest first site is not automatically best.

Part 3: any number of stages, many sites

Support any number of stages. The total number of sites K across all stages can be up to 10^5, positions are integers in [-10^9, 10^9] and build costs are in [0, 10^9]. A solution that compares every pair of sites between consecutive stages is too slow when one stage has many sites.

Constraints for all parts: every factory or stage has at least one option. An empty input returns 0.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

Clarify first. The reported question comes in parts. Pin down the model for each before coding:

  • Part 1: costs[i] lists the build options for factory i. Rows can have different lengths. Pick exactly one option per factory. Is an empty row possible? (Treat it as invalid input.) Is an empty costs list possible? (Total is 0.)
  • Parts 2-3: the factories are now stages of one production line, built in order. Each option is a site (position, build_cost) on a line. Goods move from the stage-i site to the stage-i+1 site, which costs |x - y|. Minimize build costs plus transport costs.
  • Ask about value ranges. With 10^5 options and costs up to 10^9, totals pass 32 bits. Python does not care, but say it.

Part 1: the choices are independent. Nothing links factory i to factory j, so the minimum of the sum is the sum of the minimums. Take min(row) for each row. O(K) for K total options. Do not reach for DP here; the interviewer wants to see you notice independence.

Part 2: a transit cost couples neighbours. Now a cheap site can be a bad choice if it is far from the next stage. Brute force over every combination is prod(k_i), which explodes. But the cost only couples consecutive stages, so this is a shortest path in a layered graph:

dp[i][j] = cost(i, j) + min over p of (dp[i-1][p] + |x_j - x_p|)

The answer is min(dp[last]). Computed directly, each layer costs k_{i-1} * k_i. With three stages of a few hundred sites each, this is fine and is the expected answer for part 2.

Part 3: many stages, many sites. Remove the pairwise min. Split the absolute value by side:

  • For a previous site at px <= x: dp_p + x - px = (dp_p - px) + x.
  • For a previous site at px >= x: dp_p + px - x = (dp_p + px) - x.

Sort both layers by position. Sweep left to right, keeping the running minimum of dp_p - px over previous sites at or left of x; sweep right to left with the running minimum of dp_p + px. Each new site takes the better of the two plus its own build cost. This is two pointers over two sorted lists, so a layer costs a sort plus a linear sweep, O(K log K) overall.

Invariant. After processing stage i, dp holds, for every site of stage i, the cheapest total cost of any valid chain of picks for stages 0..i that ends at that site. Each sweep sees exactly the previous sites on its side of x, and every previous site is on at least one side, so no candidate is missed. Using |x - px| on the "wrong" side is never needed because the split is exact.

Code

from __future__ import annotations

import math
from itertools import product


def min_build_cost(costs: list[list[int]]) -> int:
    """Part 1: factories are independent, so take each factory's cheapest option."""
    if any(not row for row in costs):
        raise ValueError("every factory needs at least one option")
    return sum(min(row) for row in costs)


def min_network_cost_naive(stages: list[list[tuple[int, int]]]) -> int:
    """Part 2 reference: try every pair between consecutive stages. O(sum k_i * k_{i+1})."""
    if not stages:
        return 0
    dp = [cost for _, cost in stages[0]]
    for prev, cur in zip(stages, stages[1:]):
        dp = [
            cost + min(d + abs(x - px) for d, (px, _) in zip(dp, prev))
            for x, cost in cur
        ]
    return min(dp)


def min_network_cost(stages: list[list[tuple[int, int]]]) -> int:
    """Part 3: one site per stage; pay build cost plus |x - y| between consecutive sites.

    stages[i] is a list of (position, build_cost) options for stage i.
    Runs in O(K log K) where K is the total number of options.
    """
    if not stages:
        return 0
    if any(not s for s in stages):
        raise ValueError("every stage needs at least one option")

    # dp holds (position, best total cost ending at that site), sorted by position.
    dp = sorted((x, c) for x, c in stages[0])
    for stage in stages[1:]:
        cur = sorted(stage)
        n, m = len(dp), len(cur)
        best = [math.inf] * m

        # Sweep left to right: sites at px <= x pay (d - px) + x.
        i, run = 0, math.inf
        for j, (x, _) in enumerate(cur):
            while i < n and dp[i][0] <= x:
                run = min(run, dp[i][1] - dp[i][0])
                i += 1
            best[j] = run + x

        # Sweep right to left: sites at px >= x pay (d + px) - x.
        i, run = n - 1, math.inf
        for j in range(m - 1, -1, -1):
            x = cur[j][0]
            while i >= 0 and dp[i][0] >= x:
                run = min(run, dp[i][1] + dp[i][0])
                i -= 1
            best[j] = min(best[j], run - x)

        dp = [(x, int(best[j]) + c) for j, (x, c) in enumerate(cur)]
    return min(d for _, d in dp)


if __name__ == "__main__":
    assert min_build_cost([[4, 2, 7], [5], [3, 9]]) == 10
    example = [[(0, 5), (10, 1)], [(3, 2), (9, 6)], [(4, 1)]]
    assert min_network_cost(example) == min_network_cost_naive(example) == 12
    assert min_network_cost([[(-5, 0)], [(5, 0)]]) == 10
    print("ok")

Complexity

Let K be the total number of options across all factories, and k_i the options at stage i.

  • min_build_cost: O(K) time, O(1) extra space.
  • min_network_cost_naive: O(sum k_i * k_{i+1}) time, O(max k_i) space (only the previous layer is kept).
  • min_network_cost: O(K log K) time for the per-stage sorts plus O(K) for the sweeps; O(max k_i) space. If positions arrive sorted, it is O(K).

Test cases

  • Part 1: [[4, 2, 7], [5], [3, 9]] -> 10 (2 + 5 + 3). Rows of different lengths are fine.
  • Part 1: [] -> 0. A row [] -> ValueError.
  • Part 1: one factory [[8, 8, 8]] -> 8 (ties do not matter).
  • Part 3: [[(0, 5), (10, 1)], [(3, 2), (9, 6)], [(4, 1)]] -> 12. Both 0 -> 3 -> 4 (5+3+2+1+1) and 10 -> 3 -> 4 (1+7+2+1+1) cost 12; the cheapest first site is not automatically better.
  • Part 3: one stage [[(7, 3)]] -> 3 (no transit).
  • Part 3: every stage at the same position [[(5, 1)], [(5, 1)], [(5, 1)]] -> 3 (zero transit; checks the <= / >= ties are counted).
  • Part 3: negative positions [[(-5, 0)], [(5, 0)]] -> 10.
  • Part 3: a site cheaper to build but far away loses: [[(0, 0)], [(100, 0), (1, 50)]] -> 51.
  • Randomized: for small inputs, min_network_cost == min_network_cost_naive == brute force over itertools.product.

Follow-ups

  • Return the chosen sites, not just the cost. Store, for each site, which previous site gave its best value (track the argmin alongside each running minimum). Walk back from the best final site. Same complexity, O(K) extra space.
  • Transit cost is weighted, w_i * |x - y|, or differs per stage. The split still works: keep running minima of dp_p - w_i * px and dp_p + w_i * px. A general cost such as (x - y)^2 breaks the split; it needs convex hull trick or Li Chao tree for O(K log K), or falls back to the naive DP.
  • Sites are on a 2D grid with Manhattan distance. |dx| + |dy| gives four quadrant cases. Each is a 2D dominance query: sort by x, sweep, and keep a Fenwick tree over compressed y for prefix minima. O(K log K) per layer.
  • Stages are a DAG, not a chain (a stage can feed several later stages, or the final product needs parts from several factories). If the structure is a tree, run the same DP bottom-up and add children's best values. A general DAG with shared parts is no longer a simple shortest path.
  • Budget or capacity limits (at most B total build cost, or one site can serve two stages). Budget adds a dimension to the state (dp[site][budget]) and becomes pseudo-polynomial; ask for ranges before committing.

Interviewer expectations

Mid-level (L4)

  • Solves part 1 quickly and explains why independent choices mean sum of minimums.
  • For part 2, writes the layered DP with the naive pairwise transition and gets it correct, including a single stage and empty input.
  • States time and space correctly and tests a case where the cheapest first site is not optimal.

Senior (L5)

  • Frames part 2 as a shortest path through layers and keeps only one layer of state.
  • Reaches the O(K log K) split of the absolute value for many stages, with ties handled.
  • Checks the fast version against the naive one or a brute force; can reconstruct the chosen sites.

Staff (L6+)

  • Drives the clarifications (input shape, ranges, overflow, what to return) before coding and keeps one structure that extends across all parts.
  • Explains which cost functions keep the trick working and what to use when they do not (convex hull trick, Fenwick tree for 2D).
  • Discusses the modelling choices a real factory network would add: shared parts, capacity, and how the problem stops being a chain.

Sign in to comment, vote or remix.