LeetCode/go/robot-collisions.go

package main

import (
	"cmp"
	"slices"
)

type Robot struct {
	index     int
	position  int
	health    int
	direction byte
}

func survivedRobotsHealths(positions []int, healths []int, directions string) []int {
	// construct a slice of robots instead of separate slices
	robots := make([]Robot, len(positions))
	for i := range len(positions) {
		robots[i] = Robot{
			index:     i,
			position:  positions[i],
			health:    healths[i],
			direction: directions[i],
		}
	}

	// sort by positions
	slices.SortFunc(robots, func(a, b Robot) int {
		return cmp.Compare(a.position, b.position)
	})

	var stack []*Robot
	for i := range robots {
		robot := &robots[i]
		if robot.direction == 'R' {
			stack = append(stack, robot)
			continue
		}

		for len(stack) > 0 && robot.health > 0 {
			otherRobot := stack[len(stack)-1]
			stack = stack[:len(stack)-1]

			if otherRobot.health > robot.health {
				otherRobot.health--
				robot.health = 0
				stack = append(stack, otherRobot)
			} else if otherRobot.health < robot.health {
				robot.health--
				otherRobot.health = 0
			} else {
				robot.health = 0
				otherRobot.health = 0
			}
		}
	}

	// revert to the original ordering to preserve the order of the robots
	// for the output
	slices.SortFunc(robots, func(a, b Robot) int {
		return cmp.Compare(a.index, b.index)
	})

	var hp []int
	for _, robot := range robots {
		if robot.health > 0 {
			hp = append(hp, robot.health)
		}
	}
	return hp
}