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 }