Thomas Ngo

Software Engineer

2 minute read

Problem Statement

2751

Intuition

Initially, I need to simulate the movement of the robots and handle the collisions based on their positions, healths, and directions. Robots moving towards each other will collide and the one with higher health will continue with reduced health, while the one with lower health will be removed.

Approach

  1. Initialize Robots: Create a list of robots with their positions, healths, directions, and original order.
  2. Sort by Positions: Sort the robots by their positions to handle collisions in order.
  3. Simulate Collisions: Use a stack to simulate the process of robots moving and colliding:
    • Iterate through the list and push robots onto the stack.
    • When a collision is detected (i.e., a right-moving robot meets a left-moving robot), compare their healths.
    • The robot with higher health survives with its health decremented by one, and the other robot is removed.
  4. Check Stability: Repeat the process until no more collisions occur.
  5. Sort by Original Order: Once the final state is achieved, sort the robots by their original order to return the results in the correct sequence.
  6. Return Results: Return the healths of the surviving robots.

Complexity

  • Time Complexity: The sorting step contributes \(O(n \log n)\), and the collision handling in the worst case can be \(O(n)\) for each pass. Since the number of passes depends on the collisions, in the worst case it can be \(O(n^2)\).
  • Space Complexity: The space complexity is \(O(n)\) due to the additional space required to store robot information and the stack.

Code

class Solution:
    def survivedRobotsHealths(self, positions: List[int], healths: List[int], directions: str) -> List[int]:
        robots = []
        n = len(positions)

        # Initialize robots with positions, healths, directions, and original order
        for i in range(n):
            robot = {
                'position': positions[i],
                'health': healths[i],
                'direction': -1 if directions[i] == 'L' else 1,
                'order': i
            }
            robots.append(robot)

        # Sort robots by their positions
        robots.sort(key=lambda x: x['position'])

        while True:
            stack = []
            while robots:
                robot = robots.pop()
                stack.append(robot)

                # Handle collisions
                while len(stack) >= 2 and stack[-1]['direction'] == 1 and stack[-2]['direction'] == -1:
                    robot_a = stack.pop()
                    robot_b = stack.pop()
                    if robot_a['health'] > robot_b['health']:
                        robot_a['health'] -= 1
                        stack.append(robot_a)
                    elif robot_a['health'] < robot_b['health']:
                        robot_b['health'] -= 1
                        stack.append(robot_b)

            if not stack:
                return []

            # Prepare for next pass
            robots = stack[:]
            robots.sort(key=lambda x: x['position'])

            l = 0
            direction = robots[l]['direction']
            while l < len(robots) and robots[l]['direction'] == direction:
                l += 1

            if l == len(robots):
                break

            if l >= 0 and robots[l - 1]['direction'] == -1 and robots[l]['direction'] == 1:
                break

        # Sort by original order and return the healths of surviving robots
        robots.sort(key=lambda x: x['order'])
        return [robot['health'] for robot in robots]

Editorial

Approach: Sorting & Stack

class Solution:
    def survivedRobotsHealths(
        self, positions: List[int], healths: List[int], directions: str
    ) -> List[int]:
        n = len(positions)
        indices = list(range(n))
        result = []
        stack = deque()

        # Sort indices based on their positions
        indices.sort(key=lambda x: positions[x])

        for current_index in indices:
            # Add right-moving robots to the stack
            if directions[current_index] == "R":
                stack.append(current_index)
            else:
                while stack and healths[current_index] > 0:
                    # Pop the top robot from the stack for collision check
                    top_index = stack.pop()

                    if healths[top_index] > healths[current_index]:
                        # Top robot survives, current robot is destroyed
                        healths[top_index] -= 1
                        healths[current_index] = 0
                        stack.append(top_index)
                    elif healths[top_index] < healths[current_index]:
                        # Current robot survives, top robot is destroyed
                        healths[current_index] -= 1
                        healths[top_index] = 0
                    else:
                        # Both robots are destroyed
                        healths[current_index] = 0
                        healths[top_index] = 0

        # Collect surviving robots
        for index in range(n):
            if healths[index] > 0:
                result.append(healths[index])

        return result
  • time: O(n log n)
  • space: O(n)

You may also enjoy