basis/math/combinatorics/combinatorics.factor

   1 ! Copyright (c) 2007-2010 Slava Pestov, Doug Coleman, Aaron Schaefer, John Benediktsson.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3
   4 USING: accessors arrays assocs binary-search fry kernel locals
   5 math math.order math.ranges namespaces sequences sorting ;
   6
   7 IN: math.combinatorics
   8
   9 <PRIVATE
  10
  11 : possible? ( n m -- ? )
  12     0 rot between? ; inline
  13
  14 : twiddle ( n k -- n k )
  15     2dup - dupd > [ dupd - ] when ; inline
  16
  17 PRIVATE>
  18
  19 : factorial ( n -- n! )
  20     iota 1 [ 1 + * ] reduce ;
  21
  22 : nPk ( n k -- nPk )
  23     2dup possible? [ dupd - [a,b) product ] [ 2drop 0 ] if ;
  24
  25 : nCk ( n k -- nCk )
  26     twiddle [ nPk ] keep factorial / ;
  27
  28
  29 ! Factoradic-based permutation methodology
  30
  31 <PRIVATE
  32
  33 : factoradic ( n -- factoradic )
  34     0 [ over 0 > ] [ 1 + [ /mod ] keep swap ] produce reverse 2nip ;
  35
  36 : (>permutation) ( seq n -- seq )
  37     [ '[ _ dupd >= [ 1 + ] when ] map ] keep prefix ;
  38
  39 : >permutation ( factoradic -- permutation )
  40     reverse 1 cut [ (>permutation) ] each ;
  41
  42 : permutation-indices ( n seq -- permutation )
  43     length [ factoradic ] dip 0 pad-head >permutation ;
  44
  45 PRIVATE>
  46
  47 : permutation ( n seq -- seq' )
  48     [ permutation-indices ] keep nths ;
  49
  50 : all-permutations ( seq -- seq' )
  51     [ length factorial iota ] keep
  52     '[ _ permutation ] map ;
  53
  54 : each-permutation ( seq quot -- )
  55     [ [ length factorial iota ] keep ] dip
  56     '[ _ permutation @ ] each ; inline
  57
  58 : reduce-permutations ( seq identity quot -- result )
  59     swapd each-permutation ; inline
  60
  61 : inverse-permutation ( seq -- permutation )
  62     <enum> sort-values keys ;
  63
  64
  65 ! Combinadic-based combination methodology
  66
  67 <PRIVATE
  68
  69 TUPLE: combo
  70     { seq sequence }
  71     { k integer } ;
  72
  73 C: <combo> combo
  74
  75 : choose ( combo -- nCk )
  76     [ seq>> length ] [ k>> ] bi nCk ;
  77
  78 : largest-value ( a b x -- v )
  79     dup 0 = [
  80         drop 1 - nip
  81     ] [
  82         [ iota ] 2dip '[ _ nCk _ >=< ] search nip
  83     ] if ;
  84
  85 :: next-values ( a b x -- a' b' x' v )
  86     a b x largest-value dup :> v  ! a'
  87     b 1 -                         ! b'
  88     x v b nCk -                   ! x'
  89     v ;                           ! v == a'
  90
  91 : dual-index ( m combo -- m' )
  92     choose 1 - swap - ;
  93
  94 : initial-values ( combo m -- n k m )
  95     [ [ seq>> length ] [ k>> ] bi ] dip ;
  96
  97 : combinadic ( combo m -- combinadic )
  98     initial-values [ over 0 > ] [ next-values ] produce
  99     [ 3drop ] dip ;
 100
 101 :: combination-indices ( m combo -- seq )
 102     combo m combo dual-index combinadic
 103     combo seq>> length 1 - swap [ - ] with map ;
 104
 105 : apply-combination ( m combo -- seq )
 106     [ combination-indices ] keep seq>> nths ;
 107
 108 : combinations-quot ( seq k quot -- seq quot )
 109     [ <combo> [ choose iota ] keep ] dip
 110     '[ _ apply-combination @ ] ; inline
 111
 112 PRIVATE>
 113
 114 : each-combination ( seq k quot -- )
 115     combinations-quot each ; inline
 116
 117 : map-combinations ( seq k quot -- )
 118     combinations-quot map ; inline
 119
 120 : map>assoc-combinations ( seq k quot exemplar -- )
 121     [ combinations-quot ] dip map>assoc ; inline
 122
 123 : combination ( m seq k -- seq' )
 124     <combo> apply-combination ;
 125
 126 : all-combinations ( seq k -- seq' )
 127     [ ] map-combinations ;
 128
 129 : reduce-combinations ( seq k identity quot -- result )
 130     [ -rot ] dip each-combination ; inline
 131
 132 : all-subsets ( seq -- subsets )
 133     dup length [0,b] [ all-combinations ] with map concat ;
 134
 135 <PRIVATE
 136
 137 : (selections) ( seq n -- selections )
 138     [ [ 1array ] map dup ] [ 1 - ] bi* [
 139         cartesian-product concat [ { } concat-as ] map
 140     ] with times ;
 141
 142 PRIVATE>
 143
 144 : selections ( seq n -- selections )
 145     dup 0 > [ (selections) ] [ 2drop { } ] if ;
 146
 147