Add combination support to math.combinatorics

basis/math/combinatorics/combinatorics.factor

@@ -1,7 +1,7 @@
-! Copyright (c) 2007, 2008 Slava Pestov, Doug Coleman, Aaron Schaefer.
+! Copyright (c) 2007-2009 Slava Pestov, Doug Coleman, Aaron Schaefer.
 ! See http://factorcode.org/license.txt for BSD license.
-USING: assocs kernel math math.order math.ranges mirrors
-namespaces sequences sorting fry ;
+USING: accessors assocs fry kernel locals math math.order math.ranges mirrors
+    namespaces sequences sorting ;
 IN: math.combinatorics
 
 <PRIVATE
@@ -12,14 +12,27 @@ IN: math.combinatorics
 : twiddle ( n k -- n k )
     2dup - dupd > [ dupd - ] when ; inline
 
-! See this article for explanation of the factoradic-based permutation methodology:
-! http://msdn2.microsoft.com/en-us/library/aa302371.aspx
+PRIVATE>
+
+: factorial ( n -- n! )
+    1 [ 1 + * ] reduce ;
+
+: nPk ( n k -- nPk )
+    2dup possible? [ dupd - [a,b) product ] [ 2drop 0 ] if ;
+
+: nCk ( n k -- nCk )
+    twiddle [ nPk ] keep factorial / ;
+
+
+! Factoradic-based permutation methodology
+
+<PRIVATE
 
 : factoradic ( n -- factoradic )
-    0 [ over 0 > ] [ 1+ [ /mod ] keep swap ] produce reverse 2nip ;
+    0 [ over 0 > ] [ 1 + [ /mod ] keep swap ] produce reverse 2nip ;
 
 : (>permutation) ( seq n -- seq )
-    [ '[ _ dupd >= [ 1+ ] when ] map ] keep prefix ;
+    [ '[ _ dupd >= [ 1 + ] when ] map ] keep prefix ;
 
 : >permutation ( factoradic -- permutation )
     reverse 1 cut [ (>permutation) ] each ;
@@ -29,15 +42,6 @@ IN: math.combinatorics
 
 PRIVATE>
 
-: factorial ( n -- n! )
-    1 [ 1+ * ] reduce ;
-
-: nPk ( n k -- nPk )
-    2dup possible? [ dupd - [a,b) product ] [ 2drop 0 ] if ;
-
-: nCk ( n k -- nCk )
-    twiddle [ nPk ] keep factorial / ;
-
 : permutation ( n seq -- seq )
     [ permutation-indices ] keep nths ;
 
@@ -53,3 +57,39 @@ PRIVATE>
 
 : inverse-permutation ( seq -- permutation )
     <enum> >alist sort-values keys ;
+
+
+! Combinadic-based combination methodology
+
+TUPLE: combination
+    { n integer }
+    { k integer } ;
+
+C: <combination> combination
+
+<PRIVATE
+
+: dual-index ( combination m -- x )
+    [ [ n>> ] [ k>> ] bi nCk 1 - ] dip - ;
+
+: largest-value ( a b x -- v )
+    #! TODO: use a binary search instead of find-last
+    [ [0,b) ] 2dip '[ _ nCk _ <= ] find-last nip ;
+
+:: next-values ( a b x -- a' b' x' v )
+    a b x largest-value dup :> v  ! a'
+    b 1 -                         ! b'
+    x v b nCk -                   ! x'
+    v ;                           ! v == a'
+
+: initial-values ( combination m -- a b x )
+    [ [ n>> ] [ k>> ] [ ] tri ] dip dual-index ;
+
+: combinadic ( combination m -- combinadic )
+    initial-values [ over 0 > ] [ next-values ] produce
+    [ 3drop ] dip ;
+
+PRIVATE>
+
+: combination ( m combination -- seq )
+    swap [ drop n>> 1 - ] [ combinadic ] 2bi [ - ] with map ;