Interview question

interview-problem.go

/*
We have a list of destinations (cells) and a list of agents. We want to find the closest agent-destination pair. If multiple pairs are found, we choose the agent with the lower priority. Then, the agent walks to the destination and the destination is removed from the list of destinations. The agent then becomes free.

We want to assign agents to destinations as soon as possible. We can model this problem in Go as follows:
*/
package main

import (
	"fmt"
	"sort"
	"sync/atomic"
)

// for tracking the number of available agents
var availableAgentsNum atomic.Int32

const (
	Left  Direction = -1
	Right Direction = 1
	Up    Direction = 1
	Down  Direction = -1
)

type Direction int

type Agent struct {
	id       int
	priority int
	cell     Cell
}

type Cell struct {
	x, y int
}

type Dist struct {
	length    int
	agent     *Agent
	cell      Cell
	cellIndex int // for tracking index in original cells
}

func (d Direction) toInt() int {
	return int(d)
}

// Don't send the signal if no cells remained.
func (agent *Agent) Walk(cell Cell, ch chan *Agent, sendSignal bool) {
	fmt.Println("----------------------")
	fmt.Printf("Agent %d is at (%d, %d)\n", agent.id, agent.x(), agent.y())
	fmt.Printf("Agent %d wants to go to cell (%d, %d)\n", agent.id, cell.x, cell.y)

	var xDir Direction = Left
	var yDir Direction = Down

	if agent.x() < cell.x {
		xDir = Right
	}

	if agent.y() < cell.y {
		yDir = Up
	}

	for !agent.isAt(cell) {
		if agent.x() != cell.x {
			agent.addX(xDir.toInt())
		}
		if agent.y() != cell.y {
			agent.addY(yDir.toInt())
		}
		fmt.Printf("Agent %d is at (%d, %d)\n", agent.id, agent.x(), agent.y())
	}

	if sendSignal {
		availableAgentsNum.Add(1)
		ch <- agent
	}
}

func (agent *Agent) x() int {
	return agent.cell.x
}

func (agent *Agent) y() int {
	return agent.cell.y
}

func (agent *Agent) addX(x int) {
	agent.cell.x += x
}

func (agent *Agent) addY(y int) {
	agent.cell.y += y
}

func (agent *Agent) isAt(cell Cell) bool {
	return agent.cell == cell
}

func (agent *Agent) distFrom(cell Cell) int {
	x := abs(agent.x() - cell.x)
	y := abs(agent.y() - cell.y)
	if x > y {
		return x
	}
	return y
}

func fetchAvailableAgents(ch chan *Agent) []*Agent {
	agents := []*Agent{}
	n := availableAgentsNum.Load()

	// We must wait for atleast one agent
	if n == 0 {
		agents = append(agents, <-ch)
		availableAgentsNum.Add(-1)
		// Now we can choose more agents if available
		n = availableAgentsNum.Load()
	}

	for i := int32(0); i < n; i++ {
		agents = append(agents, <-ch)
	}

	availableAgentsNum.Add(-n)

	return agents
}

func computeAllDistances(agents []*Agent, cells []Cell) []Dist {
	allDists := []Dist{}

	for _, agent := range agents {
		for index, cell := range cells {
			allDists = append(allDists,
				Dist{
					length:    agent.distFrom(cell),
					agent:     agent,
					cell:      cell,
					cellIndex: index,
				})
		}
	}

	// sort according distance and priority
	sort.Slice(allDists, func(i, j int) bool {
		d1 := allDists[i]
		d2 := allDists[j]

		return d1.length < d2.length || (d1.length == d2.length && d1.agent.priority < d2.agent.priority)
	})

	return allDists
}

// returns selected agent to walk to cell.
// Also removes the cell from the unvisited cells
func matchAgentCell(cells []Cell, ch chan *Agent) (*Agent, Cell, []Cell) {
	agents := fetchAvailableAgents(ch)
	allDists := computeAllDistances(agents, cells)

	// Desired matching is at zero index
	// Match agant to cell
	agent := allDists[0].agent
	cell := allDists[0].cell
	cellIndex := allDists[0].cellIndex

	// Remove the chosen cell by swapping to the end
	m := len(cells)
	cells[cellIndex], cells[m-1] = cells[m-1], cells[cellIndex]
	cells = cells[:m-1]

	return agent, cell, cells
}

func abs(x int) int {
	if x < 0 {
		return -x
	}
	return x
}

func sampleAgents() []*Agent {
	return []*Agent{
		{
			id:       1,
			priority: 5,
			cell: Cell{
				x: -4,
				y: 15,
			},
		}, {
			id:       2,
			priority: 4,
			cell: Cell{
				x: -20,
				y: -20,
			},
		}, {
			id:       3,
			priority: 5,
			cell: Cell{
				x: 4,
				y: 0,
			},
		},
	}
}

func sampleCells() []Cell {
	return []Cell{
		{x: 7, y: 8},
		{x: 0, y: -10},
		{x: -21, y: 16},
		{x: -10, y: -10},
	}
}

func main() {
	agents := sampleAgents()
	cells := sampleCells()

	// Agents announce availability through this channel.
	ch := make(chan *Agent)
	defer close(ch)

	n := len(agents)
	availableAgentsNum.Add(int32(n))

	// In the beginning, all agents are available.
	for i := 0; i < n; i++ {
		go func(pos int) {
			ch <- agents[pos]
		}(i)
	}

	for len(cells) != 0 {
		agent, cell, remainedCells := matchAgentCell(cells, ch)
		cells = remainedCells
		sendSignal := len(cells) != 0 // Send a signal only in case of remaining cells.
		go agent.Walk(cell, ch, sendSignal)
	}
}