factor: trim using lists

[factor.git] / extra / sorting / extras / extras.factor

diff --git a/extra/sorting/extras/extras.factor b/extra/sorting/extras/extras.factor
index 6f0ef81104af5037b66ef04d39bd7a2824bc5464..9dcb43aa390ab7b424ea987b5ee2988c25f46fe3 100644 (file)
--- a/extra/sorting/extras/extras.factor
+++ b/extra/sorting/extras/extras.factor
@@ -1,7 +1,37 @@
-USING: assocs kernel sequences sequences.private sorting ;
+USING: arrays assocs kernel kernel.private math math.order
+sequences sequences.extras sequences.private sorting ;
 IN: sorting.extras
 
 : argsort ( seq quot: ( obj1 obj2 -- <=> ) -- sortedseq )
-    [ dup length iota zip ] dip
+    [ zip-index ] dip
     [ [ first-unsafe ] bi@ ] prepose
-    sort [ 1 swap nth-unsafe ] map! ; inline
+    sort [ second-unsafe ] map! ; inline
+
+: map-sort ( ... seq quot: ( ... elt -- ... key ) -- ... sortedseq )
+    [ keep ] curry { } map>assoc
+    [ { array } declare first-unsafe ] sort-with
+    [ { array } declare second-unsafe ] map ; inline
+
+:: bisect-left ( obj seq -- i )
+    0 seq length [ 2dup < ] [
+        2dup + 2/ dup seq nth-unsafe obj before?
+        [ swap [ nip 1 + ] dip ] [ nip ] if
+    ] while drop ;
+
+:: bisect-right ( obj seq -- i )
+    0 seq length [ 2dup < ] [
+        2dup + 2/ dup seq nth-unsafe obj after?
+        [ nip ] [ swap [ nip 1 + ] dip ] if
+    ] while drop ;
+
+: insort-left ( obj seq -- seq' )
+    [ bisect-left ] 2keep swapd insert-nth ;
+
+: insort-right ( obj seq -- seq' )
+    [ bisect-right ] 2keep swapd insert-nth ;
+
+: insort-left! ( obj seq -- seq )
+    [ bisect-left ] 2keep swapd [ insert-nth! ] keep ;
+
+: insort-right! ( obj seq -- seq )
+    [ bisect-right ] 2keep swapd [ insert-nth! ] keep ;