golang argsort

argsort.go

// https://play.golang.org/p/A3niFJ582wJ

package main

import (
	"fmt"
	"sort"
)

func main() {
	fmt.Println("Hello, playground")
	
	example := []float64{1, 25, 3, 5, 4}
	fmt.Println("pre :", example)
	ia := ArgsortNew(example)
	fmt.Println("post:", example)
	fmt.Println("ia  :", ia)
	sortedExample := []float64{}
	for i := range ia {
		sortedExample = append(sortedExample, example[ia[i]])
	}
	fmt.Println("sortedExample:", sortedExample)	
}

//----------

// argsort, like in Numpy, it returns an array of indexes into an array. Note
// that the gonum version of argsort reorders the original array and returns
// indexes to reconstruct the original order.
type argsort struct {
	s    []float64 // Points to orignal array but does NOT alter it.
	inds []int     // Indexes to be returned.
}

func (a argsort) Len() int {
	return len(a.s)
}

func (a argsort) Less(i, j int) bool {
	return a.s[a.inds[i]] < a.s[a.inds[j]]
}

func (a argsort) Swap(i, j int) {
	a.inds[i], a.inds[j] = a.inds[j], a.inds[i]
}

// ArgsortNew allocates and returns an array of indexes into the source float
// array.
func ArgsortNew(src []float64) []int {
	inds := make([]int, len(src))
	for i := range src {
		inds[i] = i
	}
	Argsort(src, inds)
	return inds
}

// Argsort alters a caller-allocated array of indexes into the source float
// array. The indexes must already have values 0...n-1.
func Argsort(src []float64, inds []int) {
	if len(src) != len(inds) {
		panic("floats: length of inds does not match length of slice")
	}
	a := argsort{s: src, inds: inds}
	sort.Sort(a)
}

if you need a package that does it with any sort.Interface, here is one: https://godoc.org/github.com/mkmik/argsort