Get the indices of the array after sorting in golang

The accepted answer is good! But I needed a version of it that "didn't touch" my slice.

You can do it like this:

type sortable struct {
    nums, idxs []int
}

func (s sortable) Len() int           { return len(s.nums) }
func (s sortable) Less(i, j int) bool { return s.nums[i] < s.nums[j] }
func (s sortable) Swap(i, j int) {
    s.nums[i], s.nums[j] = s.nums[j], s.nums[i]
    s.idxs[i], s.idxs[j] = s.idxs[j], s.idxs[i]
}

func sortAndReturnIdxs(nums []int) []int {
    idxs := make([]int, len(nums))
    for i := range idxs {
        idxs[i] = i
    }

    sort.Sort(sortable{nums, idxs})

    return idxs
}

And then you can just call the function on your slice:

func main() {
    nums := []int{4, 1, 6}

    fmt.Println(sortAndReturnIdxs(nums)) // [1 0 2]
    fmt.Println(nums)                    // [1 4 6]
}

The function returns the indices, and sorts the array in place.

Playground link: https://go.dev/play/p/7tli8tNeWNt


Make a wrapper for sort.IntSlice that remembers the indexes and swaps them when it swaps the values:

type Slice struct {
    sort.IntSlice
    idx []int
}

func (s Slice) Swap(i, j int) {
    s.IntSlice.Swap(i, j)
    s.idx[i], s.idx[j] = s.idx[j], s.idx[i]
}

Playground: http://play.golang.org/p/LnSLfe-fXk.

EDIT: As DaveC mentioned in the comments, you can actually wrap around sort.Interface to create a data structure for any sortable type:

type Slice struct {
    sort.Interface
    idx []int
}

func (s Slice) Swap(i, j int) {
    s.Interface.Swap(i, j)
    s.idx[i], s.idx[j] = s.idx[j], s.idx[i]
}

Tags:

Sorting

Go