diff --git a/extra/sequences/extras/extras-tests.factor b/extra/sequences/extras/extras-tests.factor index 3efcca70a9..f4de100d97 100644 --- 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 { { } } [ { } >array ] unit-test { { 1 3 } } [ 5 iota >array ] unit-test { { 1 3 5 } } [ 6 iota >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 b7a897c3e8..2f1978950f 100644 --- 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 ;