Solution Approach

The solution implements a top-down dynamic programming approach with memoization. Here's how the implementation works:

Step 1: Sorting

robot.sort() factory.sort()

First, we sort both the robots and factories arrays in ascending order by position. This ensures we can process them sequentially and avoid crossing assignments.

Step 2: Define the DP State The function dfs(i, j) represents the minimum total distance needed to repair all robots starting from index i using factories starting from index j.

Step 3: Base Cases

if i == len(robot):  return 0 if j == len(factory):  return inf

• If i == len(robot), all robots have been assigned, so the distance is 0
• If j == len(factory), we've run out of factories but still have robots to repair, which is impossible, so we return infinity

Step 4: Recursive Transitions For each state (i, j), we have two main choices:

1. Skip the current factory:

   ans = dfs(i, j + 1)

   Move to the next factory without using the current one.

2. Use the current factory for k robots:

   t = 0 for k in range(factory[j][1]):  if i + k == len(robot):  break  t += abs(robot[i + k] - factory[j][0])  ans = min(ans, t + dfs(i + k + 1, j + 1))

   • Try assigning 1, 2, ..., up to limit robots to factory j
   • For each choice of k robots, calculate the cumulative distance t
   • The distance for robot at index i+k to reach factory j is |robot[i+k] - factory[j][0]|
   • After assigning k+1 robots to this factory, recursively solve for the remaining robots starting from i+k+1 and the next factory j+1

Step 5: Memoization The @cache decorator automatically memoizes the results of dfs(i, j), preventing redundant calculations for the same state.

Step 6: Return the Result

ans = dfs(0, 0) dfs.cache_clear() return ans

Start the recursion from the first robot and first factory, clear the cache to free memory, and return the minimum total distance.

The time complexity is O(m * n * limit) where m is the number of robots, n is the number of factories, and limit is the maximum capacity of any factory. The space complexity is O(m * n) for the memoization cache.

Example Walkthrough

Let's walk through a small example to illustrate the solution approach.

Example:

• robot = [1, 6] (2 robots at positions 1 and 6)
• factory = [[2, 1], [7, 1]] (Factory at position 2 with capacity 1, Factory at position 7 with capacity 1)

Step 1: Sort arrays

• robot = [1, 6] (already sorted)
• factory = [[2, 1], [7, 1]] (already sorted)

Step 2: Start DFS from dfs(0, 0)

We need to assign robots starting from index 0 to factories starting from index 0.

At dfs(0, 0): (Robot index 0, Factory index 0)

• Current robot at position 1, current factory at position 2 with capacity 1
• Option 1: Skip factory 0 → dfs(0, 1)
• Option 2: Send 1 robot to factory 0
  • Distance = |1 - 2| = 1
  • Continue with dfs(1, 1)
  • Total for this option: 1 + dfs(1, 1)

Let's explore Option 2 first:

At dfs(1, 1): (Robot index 1, Factory index 1)

• Current robot at position 6, current factory at position 7 with capacity 1
• Option 1: Skip factory 1 → dfs(1, 2) → returns ∞ (no more factories)
• Option 2: Send 1 robot to factory 1
  • Distance = |6 - 7| = 1
  • Continue with dfs(2, 2) → returns 0 (all robots assigned)
  • Total for this option: 1 + 0 = 1

So dfs(1, 1) = 1

Back to dfs(0, 0):

• Option 2 gives us: 1 + 1 = 2

Now let's explore Option 1:

At dfs(0, 1): (Robot index 0, Factory index 1)

• Current robot at position 1, current factory at position 7 with capacity 1
• Option 1: Skip factory 1 → dfs(0, 2) → returns ∞ (no more factories)
• Option 2: Send 1 robot to factory 1
  • Distance = |1 - 7| = 6
  • Continue with dfs(1, 2) → returns ∞ (robot at index 1 can't be assigned)
  • Total: ∞

So dfs(0, 1) = ∞

Final Result:

• dfs(0, 0) = min(∞, 2) = 2

The optimal assignment is:

• Robot at position 1 → Factory at position 2 (distance = 1)
• Robot at position 6 → Factory at position 7 (distance = 1)
• Total distance = 2

This example shows how the algorithm explores different assignments, using memoization to avoid recalculating states, and ultimately finds the minimum total distance by trying all valid combinations of robot-to-factory assignments.

Time and Space Complexity

Time Complexity: O(m^2 × n)

The time complexity analysis breaks down as follows:

• The recursive function dfs(i, j) has states defined by parameters i (robot index) and j (factory index), creating m × n unique states where m = len(robot) and n = len(factory).
• Due to memoization with @cache, each state is computed only once.
• For each state (i, j), the function iterates through up to factory[j][1] positions to assign robots to the current factory. In the worst case, this can be up to m robots (when a factory has capacity equal to or greater than the number of robots).
• Within this loop, each iteration performs O(1) operations (calculating distance and making recursive calls).
• Therefore, the work done per state is O(m), and with m × n states, the total time complexity is O(m × n × m) = O(m^2 × n).

Space Complexity: O(m × n)

The space complexity consists of:

• The memoization cache stores up to m × n states (all possible combinations of robot and factory indices).
• The recursion depth is at most O(m + n) in the worst case (processing all robots and factories).
• Since m × n dominates m + n, the overall space complexity is O(m × n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Forgetting to Sort the Arrays

One of the most common mistakes is forgetting to sort both the robot and factory arrays before applying dynamic programming. Without sorting, the greedy assignment strategy breaks down because we might assign distant robots to nearby factories, leading to suboptimal solutions.

Why it matters: The DP solution relies on the fact that once sorted, we can assign robots to factories in order without worrying about crossing assignments (where a robot further left goes to a factory further right while a robot further right goes to a factory further left).

Solution:

# Always sort both arrays at the beginning robot.sort() factory.sort()

2. Off-by-One Error in the Loop

When iterating through the number of robots to assign to a factory, it's easy to make an off-by-one error:

Incorrect:

for num_robots in range(factory_capacity + 1): # Wrong! This tries to assign capacity+1 robots  if robot_idx + num_robots >= len(robot):  break

Correct:

for num_robots in range(factory_capacity): # Assigns 0 to capacity-1 robots (total of capacity iterations)  if robot_idx + num_robots >= len(robot):  break  # When using num_robots as index offset, we're actually assigning num_robots+1 robots  cumulative_distance += abs(robot[robot_idx + num_robots] - factory_position)  min_distance = min(min_distance,  cumulative_distance + dp(robot_idx + num_robots + 1, factory_idx + 1))

3. Not Handling the "Assign Zero Robots" Case

Some implementations forget that it's valid to not assign any robots to a factory (essentially skipping it). This is already handled by having two separate options in the code, but merging them incorrectly can cause issues:

Problematic approach:

# Starting from 1 instead of considering the skip option separately for num_robots in range(1, factory_capacity + 1): # Misses the case of assigning 0 robots  ...

Correct approach:

# Option 1: Skip factory (assign 0 robots) min_distance = dp(robot_idx, factory_idx + 1)  # Option 2: Assign 1 or more robots for num_robots in range(factory_capacity):  ...

4. Integer Overflow with Large Values

When no valid assignment exists, returning a very large number is necessary. Using sys.maxsize or a large integer might cause overflow issues in some languages or when adding distances.

Solution:

# Use float('inf') for impossible cases if factory_idx == len(factory):  return float('inf')

5. Forgetting to Clear the Cache

While not affecting correctness, forgetting to clear the cache can cause memory issues in systems where the same Solution instance is reused for multiple test cases:

Solution:

result = dp(0, 0) dp.cache_clear() # Important for memory management return result