floats: add ArgsortStable function
MODIFY
floats/floats.go
floats/floats.go
@@ -74,7 +74,7 @@ func AddScaledTo(dst, y []float64, alpha float64, s []float64) []float64 {
}
// argsort is a helper that implements sort.Interface, as used by
-// Argsort.
+// Argsort and ArgsortStable.
type argsort struct {
s []float64
inds []int
@@ -110,6 +110,24 @@ func Argsort(dst []float64, inds []int) {
sort.Sort(a)
}
+// ArgsortStable sorts the elements of dst while tracking their original order and
+// keeping the original order of equal elements. At the conclusion of ArgsortStable,
+// dst will contain the original elements of dst but sorted in increasing order,
+// and inds will contain the original position of the elements in the slice such
+// that dst[i] = origDst[inds[i]].
+// It panics if the argument lengths do not match.
+func ArgsortStable(dst []float64, inds []int) {
+ if len(dst) != len(inds) {
+ panic(badDstLength)
+ }
+ for i := range dst {
+ inds[i] = i
+ }
+
+ a := argsort{s: dst, inds: inds}
+ sort.Stable(a)
+}
+
// Count applies the function f to every element of s and returns the number
// of times the function returned true.
func Count(f func(float64) bool, s []float64) int {

MODIFY
floats/floats_test.go
floats/floats_test.go
@@ -7,6 +7,7 @@ package floats
import (
"fmt"
"math"
+ "sort"
"strconv"
"testing"
@@ -180,6 +181,29 @@ func TestArgsort(t *testing.T) {
}
}
+func TestArgsortStable(t *testing.T) {
+ for i := 1; i <= 100; i++ {
+ data := make([]float64, i+2)
+ data[0] = 2
+ data[len(data)-1] = 2
+ for j := 1; j <= i; j++ {
+ data[j] = 1
+ }
+ idx := make([]int, len(data))
+ ArgsortStable(data, idx)
+ if !sort.Float64sAreSorted(data) {
+ t.Errorf("unexpected data sort order for case %d", i)
+ }
+ if !sort.IntsAreSorted(idx[:i]) {
+ t.Errorf("unexpected index sort order for case %d", i)
+ }
+ }
+
+ if !Panics(func() { ArgsortStable([]float64{1, 2, 3}, []int{0, 0}) }) {
+ t.Error("does not panic if lengths do not match")
+ }
+}
+
func TestCount(t *testing.T) {
t.Parallel()
s := []float64{3, 4, 1, 7, 5}

floats: add ArgsortStable function