LeetCode/go/relative-sort-array.go

package main

import (
	"cmp"
	"slices"
)

func relativeSortArray(arr1 []int, arr2 []int) []int {
	// remap the keys
	keys := make(map[int]int, len(arr2))
	for i, x := range arr2 {
		keys[x] = i
	}

	// sort with the given keys
	slices.SortFunc(arr1, func(x, y int) int {
		xi, xFound := keys[x]
		yi, yFound := keys[y]

		if !xFound {
			xi = 1000 + x
		}
		if !yFound {
			yi = 1000 + y
		}

		return cmp.Compare(xi, yi)
	})

	return arr1
}