sequences.extras: fixes, add "arg-max" and "arg-min".

Change supremum-by/infimum-by to return first largest or smallest element.

diff --git a/extra/sequences/extras/extras-tests.factor b/extra/sequences/extras/extras-tests.factor
index 3efcca70a997a6becbde8627f4bdaa60487e8340..f4de100d9785459dc9355e82bba35449577b9690 100644 (file)
--- a/extra/sequences/extras/extras-tests.factor
+++ b/extra/sequences/extras/extras-tests.factor
@@ -5,6 +5,8 @@ IN: sequences.extras.tests
 
 [ 4 ] [ 5 iota [ ] supremum-by ] unit-test
 [ 0 ] [ 5 iota [ ] infimum-by ] unit-test
+{ "bar" } [ { "bar" "baz" "qux" } [ length ] supremum-by ] unit-test
+{ "bar" } [ { "bar" "baz" "qux" } [ length ] infimum-by ] unit-test
 [ { "foo" } ] [ { { "foo" } { "bar" } } [ first ] supremum-by ] unit-test
 [ { "bar" } ] [ { { "foo" } { "bar" } } [ first ] infimum-by ] unit-test
 
@@ -118,3 +120,6 @@ IN: sequences.extras.tests
 { { } } [ { } <odds> >array ] unit-test
 { { 1 3 } } [ 5 iota <odds> >array ] unit-test
 { { 1 3 5 } } [ 6 iota <odds> >array ] unit-test
+
+{ 1 } [ { 1 7 3 7 6 3 7 } arg-max ] unit-test
+{ 0 } [ { 1 7 3 7 6 3 7 } arg-min ] unit-test

diff --git a/extra/sequences/extras/extras.factor b/extra/sequences/extras/extras.factor
index b7a897c3e856755ba014f4e76cc3c2eb54c3faa9..2f1978950f296df759e41de264743c44550b04bd 100644 (file)
--- a/extra/sequences/extras/extras.factor
+++ b/extra/sequences/extras/extras.factor
@@ -1,5 +1,5 @@
-USING: accessors arrays grouping kernel locals math math.order
-math.ranges sequences sequences.private splitting ;
+USING: accessors arrays assocs grouping kernel locals math
+math.order math.ranges sequences sequences.private splitting ;
 FROM: sequences => change-nth ;
 IN: sequences.extras
 
@@ -35,13 +35,13 @@ IN: sequences.extras
 
 : supremum-by ( seq quot: ( ... elt -- ... x ) -- elt )
     [ [ first dup ] dip call ] 2keep [
-        dupd call pick dupd max over =
+        dupd call pick dupd after?
         [ [ 2drop ] 2dip ] [ 2drop ] if
     ] curry 1 each-from drop ; inline
 
 : infimum-by ( seq quot: ( ... elt -- ... x ) -- elt )
     [ [ first dup ] dip call ] 2keep [
-        dupd call pick dupd min over =
+        dupd call pick dupd before?
         [ [ 2drop ] 2dip ] [ 2drop ] if
     ] curry 1 each-from drop ; inline
 
@@ -265,3 +265,9 @@ INSTANCE: odds immutable-sequence
 
 : until-empty ( seq quot -- )
     [ dup empty? ] swap until drop ; inline
+
+: arg-max ( seq -- n )
+    dup length iota zip [ first-unsafe ] supremum-by second ;
+
+: arg-min ( seq -- n )
+    dup length iota zip [ first-unsafe ] infimum-by second ;