From 7ad089b68280fb0b493fe4e4f32043b36cea6f9f Mon Sep 17 00:00:00 2001
From: Doug Coleman <doug.coleman@gmail.com>
Date: Sun, 22 Apr 2012 23:41:13 -0700
Subject: [PATCH] math.statistics: Generalize kth-smallest into kth-object, add
 kth-largest. Add count-relative and minmax-relative for counting the number
 of lt, eq, gt elements, and getting the max lt and min gt in one pass.

---
 basis/math/statistics/statistics-tests.factor |  6 ++++
 basis/math/statistics/statistics.factor       | 30 ++++++++++++++++---
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/basis/math/statistics/statistics-tests.factor b/basis/math/statistics/statistics-tests.factor
index 67323f0a41..029c25fec1 100644
--- a/basis/math/statistics/statistics-tests.factor
+++ b/basis/math/statistics/statistics-tests.factor
@@ -13,6 +13,12 @@ IN: math.statistics.tests
 [ 2 ] [ { 1 2 3 } median ] unit-test
 [ 5/2 ] [ { 1 2 3 4 } median ] unit-test
 
+{ 1 } [ { 1 2 3 4 } 0 kth-smallest ] unit-test
+{ 3 } [ { 1 2 3 4 } 2 kth-smallest ] unit-test
+
+{ 4 } [ { 1 2 3 4 } 0 kth-largest ] unit-test
+{ 2 } [ { 1 2 3 4 } 2 kth-largest ] unit-test
+
 [ 1 ] [ { 1 } mode ] unit-test
 [ 3 ] [ { 1 2 3 3 3 4 5 6 76 7 2 21 1 3 3 3 } mode ] unit-test
 
diff --git a/basis/math/statistics/statistics.factor b/basis/math/statistics/statistics.factor
index d0f203f5c7..4c9204bf80 100644
--- a/basis/math/statistics/statistics.factor
+++ b/basis/math/statistics/statistics.factor
@@ -2,7 +2,7 @@
 ! See http://factorcode.org/license.txt for BSD license.
 USING: assocs combinators generalizations kernel locals math
 math.functions math.order math.vectors sequences
-sequences.private sorting ;
+sequences.private sorting fry ;
 IN: math.statistics
 
 : mean ( seq -- x )
@@ -14,7 +14,7 @@ IN: math.statistics
 : harmonic-mean ( seq -- x )
     [ recip ] map-sum recip ;
 
-:: kth-smallest ( seq k -- elt )
+:: kth-object ( seq k quot: ( x y -- ? ) -- elt )
     #! Wirth's method, Algorithm's + Data structues = Programs p. 84
     #! The algorithm modifiers seq, so we clone it
     seq clone :> seq
@@ -30,8 +30,8 @@ IN: math.statistics
         m j!
         [ i j <= ]
         [
-            [ i seq nth-unsafe x < ] [ i 1 + i! ] while
-            [ x j seq nth-unsafe < ] [ j 1 - j! ] while
+            [ i seq nth-unsafe x quot call ] [ i 1 + i! ] while
+            [ x j seq nth-unsafe quot call ] [ j 1 - j! ] while
             i j <= [
                 i j seq exchange-unsafe
                 i 1 + i!
@@ -44,6 +44,28 @@ IN: math.statistics
     ] while
     k seq nth ; inline
 
+: kth-smallest ( seq k -- elt ) [ < ] kth-object ;
+    
+: kth-largest ( seq k -- elt ) [ > ] kth-object ;
+
+: count-relative ( seq k -- lt eq gt )
+    [ 0 0 0 ] 2dip '[
+        _ <=> {
+            { +lt+ [ [ 1 + ] 2dip ] }
+            { +gt+ [ 1 + ] }
+            { +eq+ [ [ 1 + ] dip ] }
+        } case
+    ] each ;
+
+: minmax-relative ( seq k -- lt eq gt lt-max gt-min )
+    [ 0 0 0 -1/0. 1/0. ] 2dip '[
+        dup _ <=> {
+            { +lt+ [ [ 1 + ] 5 ndip '[ _ max ] dip ] }
+            { +gt+ [ [ 1 + ] 3dip min ] }
+            { +eq+ [ [ 1 + ] 4dip drop ] }
+        } case
+    ] each ;
+    
 : lower-median ( seq -- elt )
     [ ] [ ] [ length odd? ] tri
     [ midpoint@ ] [ midpoint@ 1 - ] if kth-smallest ;