Add arg_sort and test.

google-research · Sep 24, 2023 · f456031 · f456031
1 parent 2c4b5cd
commit f456031
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/lib/sort.dx b/lib/sort.dx
@@ -63,6 +63,11 @@ def sort(xs: n=>a) -> n=>a given (n|Ix, a|Ord) =
   AsList(_, r) = reduce mempty mcombine xlists
   unsafe_cast_table(to=n, r)
 
+def arg_sort(xs: n=>a) -> n=>n given (n|Ix|Ord, a|Ord) =
+  indexed_table = for i. (xs[i], i)
+  sorted = sort indexed_table
+  map snd sorted
+
 def (+|)(i:n, delta:Nat) -> n  given (n|Ix) =
   i' = ordinal i + delta
   from_ordinal $ select (i' >= size n) (size n -| 1) i'

diff --git a/tests/sort-tests.dx b/tests/sort-tests.dx
@@ -54,3 +54,9 @@ import sort
 
 :p is_sorted $ sort ["Charlie", "Alice", "Bob", "Aaron"]
 > True
+
+'### Test argsort
+xs = ['1', '2', '3', '4', '9', '8', '7', '6']
+ixs = arg_sort xs
+sort xs == for i. xs[ixs[i]]
+> True