LeetCode Problem 2731

Difficulty: 🟡 Medium
Topics: Array, Brainteaser, Sorting, Prefix Sum

Solution

LeetCode 2731 - Movement of Robots

Problem Understanding

The problem gives us a list of robots positioned on an infinite number line. Each robot starts at a unique coordinate from the array nums, and each robot has an associated movement direction from the string s.

A robot moves exactly one unit per second. If the direction is 'L', the robot moves left, meaning its coordinate decreases over time. If the direction is 'R', the robot moves right, meaning its coordinate increases over time.

The tricky part of the problem is collision handling. Whenever two robots occupy the same position at the same moment, they instantly reverse directions. This can happen many times during the simulation.

After exactly d seconds, we must compute the sum of distances between every pair of robots. For two robots at positions x and y, their contribution is |x - y|. Since the answer can become very large, we return the result modulo 10^9 + 7.

The constraints are large:

• Up to 10^5 robots
• Coordinates can be as large as ±2 * 10^9
• Time d can also be as large as 10^9

These constraints immediately rule out any direct second-by-second simulation. Even an O(n^2) pairwise distance calculation becomes too slow when n = 10^5.

An important observation is that collisions only swap robot identities. Physically, if two identical robots collide and reverse directions, the final set of positions is exactly the same as if the robots had simply passed through each other without interacting.

This insight transforms the problem completely.

Some important edge cases include:

• d = 0, where robots never move
• All robots moving in the same direction, meaning no collisions occur
• Multiple collisions between robots
• Very large coordinates and large values of d, requiring careful handling of integer overflow
• Negative coordinates after movement

The problem guarantees that initial positions are unique, which simplifies reasoning about robot ordering.

Approaches

Brute Force Approach

A straightforward simulation would move every robot one unit at a time for d seconds. At every second, we would detect collisions and reverse directions accordingly.

After all movements complete, we would compute the distance between every pair of robots using nested loops.

This approach is correct because it directly follows the problem statement. However, it is far too slow.

If we simulate second by second, the time complexity becomes O(d * n) just for movement. Collision detection could make it even worse. Then computing pairwise distances requires another O(n^2) operations.

Given that both n and d can reach 10^5 and 10^9 respectively, this solution is infeasible.

Optimal Approach

The key observation is that collisions between identical robots are equivalent to robots passing through each other.

Suppose two robots moving toward each other collide and reverse directions. If we ignore their identities, the final positions are identical to the scenario where they simply continue moving through each other.

Therefore, we do not need to simulate collisions at all.

We can compute each robot's final position independently:

• If direction is 'R', final position is nums[i] + d
• If direction is 'L', final position is nums[i] - d

Once we obtain all final positions, the problem reduces to:

Compute the sum of pairwise absolute differences in an array.

A naive pairwise calculation would still take O(n^2), so we optimize further.

If the positions are sorted, then for each position positions[i], all previous positions are smaller or equal. The contribution of positions[i] to the total distance is:

$$positions[i] \times i - \text{sum of previous positions}$$

We maintain a running prefix sum while iterating through the sorted array.

This gives an O(n log n) solution due to sorting.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(d * n + n²)	O(n)	Simulates movement and collisions directly
Optimal	O(n log n)	O(n)	Uses collision equivalence and prefix sums

Algorithm Walkthrough

1. Create a new array to store final robot positions.

For every robot:

• If its direction is 'R', add d to its coordinate.
• If its direction is 'L', subtract d from its coordinate.

This works because collisions can be ignored mathematically. 2. Sort the final positions.

Sorting is essential because absolute differences become easier to compute in sorted order. Once sorted, every earlier element is less than or equal to the current element. 3. Initialize:

• prefix_sum = 0
• answer = 0

The prefix sum stores the sum of all previously processed positions. 4. Iterate through the sorted positions.

For each index i and value position:

• The contribution of this position against all earlier positions is:

$$position \times i - prefix_sum$$

This formula works because:

$$(position - positions[0]) + (position - positions[1]) + \dots$$

expands into:

$$position \times i - \text{sum of previous positions}$$

1. Add the contribution to the answer.
2. Update the prefix sum by adding the current position.
3. Return the answer modulo 10^9 + 7.

Why it works

The correctness relies on two properties.

First, robot collisions are equivalent to robots passing through each other. Since robots are indistinguishable except for identity, swapping directions produces the same set of final coordinates as uninterrupted motion.

Second, sorting allows us to remove the absolute value operation. In sorted order:

$$positions[i] \ge positions[j]$$

for all j < i, so:

$$|positions[i] - positions[j]| = positions[i] - positions[j]$$

This enables efficient prefix sum computation of all pairwise distances.

Python Solution

from typing import List

class Solution:
    def sumDistance(self, nums: List[int], s: str, d: int) -> int:
        MOD = 10**9 + 7
        
        final_positions = []
        
        for i, position in enumerate(nums):
            if s[i] == 'R':
                final_positions.append(position + d)
            else:
                final_positions.append(position - d)
        
        final_positions.sort()
        
        prefix_sum = 0
        answer = 0
        
        for i, position in enumerate(final_positions):
            contribution = position * i - prefix_sum
            answer = (answer + contribution) % MOD
            prefix_sum += position
        
        return answer % MOD

The implementation begins by constructing the array of final robot positions. Instead of simulating collisions, each robot simply moves directly according to its direction for d seconds.

After generating all final coordinates, the array is sorted. This ordering is important because it lets us compute pairwise distances without absolute values.

The loop then processes each position one at a time. The variable prefix_sum stores the sum of all earlier positions. Using the formula:

$$position \times i - prefix_sum$$

we efficiently compute the contribution of the current position against all previous positions in constant time.

The modulo operation is applied during accumulation to prevent excessively large intermediate values.

Go Solution

package main

import (
	"sort"
)

func sumDistance(nums []int, s string, d int) int {
	const MOD int64 = 1e9 + 7

	n := len(nums)
	finalPositions := make([]int64, n)

	for i := 0; i < n; i++ {
		if s[i] == 'R' {
			finalPositions[i] = int64(nums[i] + d)
		} else {
			finalPositions[i] = int64(nums[i] - d)
		}
	}

	sort.Slice(finalPositions, func(i, j int) bool {
		return finalPositions[i] < finalPositions[j]
	})

	var prefixSum int64 = 0
	var answer int64 = 0

	for i, position := range finalPositions {
		contribution := position*int64(i) - prefixSum
		answer = (answer + contribution) % MOD
		prefixSum += position
	}

	answer %= MOD

	if answer < 0 {
		answer += MOD
	}

	return int(answer)
}

The Go implementation follows the same algorithm as the Python version.

One important difference is integer handling. Since coordinates and intermediate sums can become very large, the implementation uses int64 throughout the computation to avoid overflow.

The positions are stored in an int64 slice before sorting. Go's sort.Slice function is used for sorting custom numeric slices.

Unlike Python, Go's modulo operator can produce negative values, so the final result is adjusted if necessary.

Worked Examples

Example 1

Input:
nums = [-2,0,2]
s = "RLL"
d = 3

Step 1: Compute Final Positions

Robot	Initial	Direction	Final
0	-2	R	1
1	0	L	-3
2	2	L	-1

Final positions:

[1, -3, -1]

Step 2: Sort Positions

[-3, -1, 1]

Step 3: Prefix Sum Computation

i	position	prefix_sum before	contribution	answer
0	-3	0	(-3 * 0) - 0 = 0	0
1	-1	-3	(-1 * 1) - (-3) = 2	2
2	1	-4	(1 * 2) - (-4) = 6	8

Final answer:

8

Example 2

Input:
nums = [1,0]
s = "RL"
d = 2

Step 1: Compute Final Positions

Robot	Initial	Direction	Final
0	1	R	3
1	0	L	-2

Final positions:

[3, -2]

Step 2: Sort Positions

[-2, 3]

Step 3: Prefix Sum Computation

i	position	prefix_sum before	contribution	answer
0	-2	0	0	0
1	3	-2	(3 * 1) - (-2) = 5	5

Final answer:

5

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log n)	Sorting dominates the runtime
Space	O(n)	Stores the transformed final positions

The algorithm first computes final positions in linear time. Sorting requires O(n log n) time, which becomes the dominant cost.

The prefix sum traversal is linear. Therefore, the total runtime is:

$$O(n \log n)$$

Additional space is required for the array of final positions, giving O(n) auxiliary space usage.

Test Cases

from typing import List

class Solution:
    def sumDistance(self, nums: List[int], s: str, d: int) -> int:
        MOD = 10**9 + 7
        
        final_positions = []
        
        for i, position in enumerate(nums):
            if s[i] == 'R':
                final_positions.append(position + d)
            else:
                final_positions.append(position - d)
        
        final_positions.sort()
        
        prefix_sum = 0
        answer = 0
        
        for i, position in enumerate(final_positions):
            contribution = position * i - prefix_sum
            answer = (answer + contribution) % MOD
            prefix_sum += position
        
        return answer % MOD

solution = Solution()

assert solution.sumDistance([-2,0,2], "RLL", 3) == 8
# Provided example with collisions

assert solution.sumDistance([1,0], "RL", 2) == 5
# Provided example with opposite movement

assert solution.sumDistance([0,10], "RR", 5) == 10
# No collisions, same direction

assert solution.sumDistance([0,1], "RL", 0) == 1
# d = 0, robots do not move

assert solution.sumDistance([-1000000000,1000000000], "LR", 1000000000) == 0
# Robots meet symmetrically

assert solution.sumDistance([1,2,3], "LLL", 1) == 4
# All robots move left

assert solution.sumDistance([1,2,3], "RRR", 1) == 4
# All robots move right

assert solution.sumDistance([-5,0,5], "RLR", 10) == 60
# Large movement distance

assert solution.sumDistance([0,2,4,6], "RLRL", 3) == 22
# Multiple crossing paths

assert solution.sumDistance([-2,-1,0,1,2], "RLRLR", 5) == 62
# Mixed directions with many effective collisions

Test	Why
[-2,0,2], "RLL", 3	Validates collision equivalence
[1,0], "RL", 2	Simple two-robot interaction
[0,10], "RR", 5	No collisions case
[0,1], "RL", 0	Zero movement edge case
[-1e9,1e9], "LR", 1e9	Large coordinate handling
[1,2,3], "LLL", 1	Uniform left movement
[1,2,3], "RRR", 1	Uniform right movement
[-5,0,5], "RLR", 10	Large displacement values
[0,2,4,6], "RLRL", 3	Multiple path crossings
[-2,-1,0,1,2], "RLRLR", 5	Stress test with mixed motion

Edge Cases

One important edge case is when d = 0. In this situation, no robot moves at all, so the answer must be computed directly from the initial positions. A buggy implementation might still attempt unnecessary transformations or collision logic. This implementation handles the case naturally because adding or subtracting zero leaves all positions unchanged.

Another important edge case occurs when many robots collide repeatedly. A direct simulation could become extremely complicated and inefficient. The mathematical observation that collisions are equivalent to robots passing through each other completely eliminates this complexity. The implementation never explicitly processes collisions.

Large coordinates and large movement distances are also important. Since values can reach billions, pairwise sums can become extremely large. In Python this is safe because integers are unbounded, but in Go we must explicitly use int64 to avoid overflow. The modulo operation ensures the final result stays within the required bounds.

A final subtle case involves negative positions after movement. Since robots may move left from already negative coordinates, final positions can become very small negative numbers. Sorting and prefix sum logic still work correctly because the algorithm depends only on relative ordering, not positivity.