smp.go

smp.go

func solve(residents, hospitals [][]int) ([][2]int, error) {
	if len(residents) != len(hospitals) {
		return nil, errors.New("smp: sets must have equal sizes")
	}

	size := len(residents)
	freeResidents := makeRange(size)

	engaged := make(map[int]int) // key: hospital, value: resident

	for len(freeResidents) > 0 {
		// get next free resident
		freeResidentIndex, updatedFreeResidents := popFront(freeResidents)
		freeResidents = updatedFreeResidents

		// get highest-ranked candidate for which resident has not yet proposed
		hospitalCandidate, availableProposals := popFront(residents[freeResidentIndex])
		residents[freeResidentIndex] = availableProposals

		if val, ok := engaged[hospitalCandidate]; ok {
			// if already engaged, check if current resident is higher ranked on hospital's list
			if indexOf(hospitals[hospitalCandidate], freeResidentIndex) > indexOf(hospitals[hospitalCandidate], val) {
				engaged[hospitalCandidate] = freeResidentIndex
				freeResidents = append(freeResidents, val)
			} else {
				freeResidents = append(freeResidents, freeResidentIndex)
			}
		} else {
			// engage hospital
			engaged[hospitalCandidate] = freeResidentIndex
		}
	}

	// build matched pairs
	matched := [][2]int{}
	for hospital, resident := range engaged {
		matched = append(matched, [2]int{resident, hospital})
	}

	return matched, nil
}

func main() {
	residents := [][]int{
		{3, 2, 1, 0},
		{2, 3, 1, 0},
		{0, 1, 2, 3},
		{2, 1, 3, 0}}

	hospitals := [][]int{
		{2, 1, 3, 0},
		{2, 0, 3, 1},
		{0, 1, 2, 3},
		{2, 1, 3, 0}}

	matched, _ := solve(residents, hospitals)

	for _, v := range matched {
		fmt.Printf("Resident %d <--> Hospital %d\n", v[0], v[1])
	}
}

/*
	Helper functions
*/

func makeRange(n int) []int {
	a := make([]int, n)
	for i := 0; i < n; i++ {
		a[i] = i
	}
	return a
}

func popBack(arr []int) (int, []int) {
	x, arr := arr[len(arr)-1], arr[:len(arr)-1]
	return x, arr
}

func popFront(arr []int) (int, []int) {
	if len(arr) < 2 {
		return arr[0], nil
	}
	x, arr := arr[0], arr[1:]
	return x, arr
}

func indexOf(arr []int, element int) int {
	for k, v := range arr {
		if v == element {
			return k
		}
	}
	return -1
}