Skip to content

2751. Robot Collisions 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
struct Robot {
  int index;
  int position;
  int health;
  char direction;
};

class Solution {
 public:
  vector<int> survivedRobotsHealths(vector<int>& positions,
                                    vector<int>& healths, string directions) {
    vector<int> ans;
    vector<Robot> robots;
    vector<Robot> stack;  // the runnnig robots

    for (int i = 0; i < positions.size(); ++i)
      robots.push_back(Robot{i, positions[i], healths[i], directions[i]});

    ranges::sort(robots, ranges::less{},
                 [](const Robot& robot) { return robot.position; });

    for (Robot& robot : robots) {
      if (robot.direction == 'R') {
        stack.push_back(robot);
        continue;
      }
      // Collide with robots going right if any.
      while (!stack.empty() && stack.back().direction == 'R' &&
             robot.health > 0) {
        if (stack.back().health == robot.health) {
          stack.pop_back();
          robot.health = 0;
        } else if (stack.back().health < robot.health) {
          stack.pop_back();
          robot.health -= 1;
        } else {  // stack.back().health > robot.health
          stack.back().health -= 1;
          robot.health = 0;
        }
      }
      if (robot.health > 0)
        stack.push_back(robot);
    }

    ranges::sort(stack, ranges::less{},
                 [](const Robot& robot) { return robot.index; });

    for (const Robot& robot : stack)
      ans.push_back(robot.health);

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class Robot {
  public int index;
  public int position;
  public int health;
  public char direction;
  public Robot(int index, int position, int health, char direction) {
    this.index = index;
    this.position = position;
    this.health = health;
    this.direction = direction;
  }
}

class Solution {
  public List<Integer> survivedRobotsHealths(int[] positions, int[] healths, String directions) {
    List<Integer> ans = new ArrayList<>();
    Robot[] robots = new Robot[positions.length];
    List<Robot> stack = new ArrayList<>(); // running robots

    for (int i = 0; i < positions.length; ++i)
      robots[i] = new Robot(i, positions[i], healths[i], directions.charAt(i));

    Arrays.sort(robots, (a, b) -> Integer.compare(a.position, b.position));

    for (Robot robot : robots) {
      if (robot.direction == 'R') {
        stack.add(robot);
        continue;
      }
      // Collide with robots going right if any.
      while (!stack.isEmpty() && stack.get(stack.size() - 1).direction == 'R' && robot.health > 0) {
        if (stack.get(stack.size() - 1).health == robot.health) {
          stack.remove(stack.size() - 1);
          robot.health = 0;
        } else if (stack.get(stack.size() - 1).health < robot.health) {
          stack.remove(stack.size() - 1);
          robot.health -= 1;
        } else { // stack[-1].health > robot.health
          stack.get(stack.size() - 1).health -= 1;
          robot.health = 0;
        }
      }
      if (robot.health > 0)
        stack.add(robot);
    }

    stack.sort((a, b) -> Integer.compare(a.index, b.index));

    for (Robot robot : stack)
      ans.add(robot.health);

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
from dataclasses import dataclass


@dataclass
class Robot:
  index: int
  position: int
  health: int
  direction: str


class Solution:
  def survivedRobotsHealths(
      self,
      positions: list[int],
      healths: list[int],
      directions: str,
  ) -> list[int]:
    robots = sorted([Robot(index, position, health, direction)
                     for index, (position, health, direction) in
                     enumerate(zip(positions, healths, directions))],
                    key=lambda x: x.position)
    stack: list[Robot] = []  # running robots

    for robot in robots:
      if robot.direction == 'R':
        stack.append(robot)
        continue
      # Collide with robots going right if any.
      while stack and stack[-1].direction == 'R' and robot.health > 0:
        if stack[-1].health == robot.health:
          stack.pop()
          robot.health = 0
        elif stack[-1].health < robot.health:
          stack.pop()
          robot.health -= 1
        else:  # stack[-1].health > robot.health
          stack[-1].health -= 1
          robot.health = 0
      if robot.health > 0:
        stack.append(robot)

    stack.sort(key=lambda robot: robot.index)
    return [robot.health for robot in stack]