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 classes.tuple
   5 combinators fry hints kernel kernel.private locals math
   6 math.functions math.order math.ranges namespaces sequences
   7 sequences.private sorting strings vectors ;
   8 IN: math.combinatorics
   9
  10 <PRIVATE
  11
  12 ! Specialized version of nths-unsafe for performance
  13 : (nths-unsafe) ( indices seq -- seq' )
  14     [ { array } declare ] dip
  15     [ [ nth-unsafe ] curry ] keep map-as ; inline
  16 GENERIC: nths-unsafe ( indices seq -- seq' )
  17 M: string nths-unsafe (nths-unsafe) ;
  18 M: array nths-unsafe (nths-unsafe) ;
  19 M: vector nths-unsafe (nths-unsafe) ;
  20 M: iota nths-unsafe (nths-unsafe) ;
  21 M: object nths-unsafe (nths-unsafe) ;
  22
  23 : possible? ( n m -- ? )
  24     0 rot between? ; inline
  25
  26 : twiddle ( n k -- n k )
  27     2dup - dupd > [ dupd - ] when ; inline
  28
  29 PRIVATE>
  30
  31 : factorial ( n -- n! )
  32     dup 1 > [ [1,b] product ] [ drop 1 ] if ;
  33
  34 : nPk ( n k -- nPk )
  35     2dup possible? [ dupd - [a,b) product ] [ 2drop 0 ] if ;
  36
  37 : nCk ( n k -- nCk )
  38     twiddle [ nPk ] keep factorial /i ;
  39
  40
  41 ! Factoradic-based permutation methodology
  42
  43 <PRIVATE
  44
  45 : factoradic ( n -- factoradic )
  46     0 [ over 0 > ] [ 1 + [ /mod ] keep swap ] produce reverse! 2nip ;
  47
  48 : bump-indices ( seq n -- )
  49     '[ dup _ >= [ 1 + ] when ] map! drop ; inline
  50
  51 : (>permutation) ( seq n index -- seq )
  52     swap [ dupd head-slice ] dip bump-indices ;
  53
  54 : >permutation ( factoradic -- permutation )
  55     reverse! dup [ (>permutation) ] each-index reverse! ;
  56
  57 : permutation-indices ( n seq -- permutation )
  58     length [ factoradic ] dip 0 pad-head >permutation ;
  59
  60 : permutation-iota ( seq -- <iota> )
  61     length factorial <iota> ; inline
  62
  63 PRIVATE>
  64
  65 : permutation ( n seq -- seq' )
  66     [ permutation-indices ] keep nths-unsafe ;
  67
  68 TUPLE: permutations length seq ;
  69
  70 : <permutations> ( seq -- permutations )
  71     [ length factorial ] keep permutations boa ;
  72
  73 M: permutations length length>> ; inline
  74 M: permutations nth-unsafe seq>> permutation ;
  75 M: permutations hashcode* tuple-hashcode ;
  76
  77 INSTANCE: permutations immutable-sequence
  78
  79 TUPLE: k-permutations length skip k seq ;
  80
  81 :: <k-permutations> ( seq k -- permutations )
  82     seq length :> n
  83     n k nPk :> len
  84     {
  85         { [ len k [ zero? ] either? ] [ { } ] }
  86         { [ n k = ] [ seq <permutations> ] }
  87         [ len n factorial over /i k seq k-permutations boa ]
  88     } cond ;
  89
  90 M: k-permutations length length>> ; inline
  91 M: k-permutations nth-unsafe
  92     [ skip>> * ]
  93     [ seq>> [ permutation-indices ] keep ]
  94     [ k>> swap [ head ] dip nths-unsafe ] tri ;
  95 M: k-permutations hashcode* tuple-hashcode ;
  96
  97 INSTANCE: k-permutations immutable-sequence
  98
  99 DEFER: next-permutation
 100
 101 <PRIVATE
 102
 103 : permutations-quot ( seq quot -- seq quot' )
 104     [ [ permutation-iota ] [ length <iota> >array ] [ ] tri ] dip
 105     '[ drop _ [ _ nths-unsafe @ ] keep next-permutation drop ] ; inline
 106
 107 PRIVATE>
 108
 109 : each-permutation ( ... seq quot: ( ... elt -- ... ) -- ... )
 110     permutations-quot each ; inline
 111
 112 : map-permutations ( ... seq quot: ( ... elt -- ... newelt ) -- ... newseq )
 113     permutations-quot map ; inline
 114
 115 : filter-permutations ( ... seq quot: ( ... elt -- ... ? ) -- ... newseq )
 116     selector [ each-permutation ] dip ; inline
 117
 118 : all-permutations ( seq -- seq' )
 119     [ ] map-permutations ;
 120
 121 : all-permutations? ( ... seq quot: ( ... elt -- ... ? ) -- ... ? )
 122     permutations-quot all? ; inline
 123
 124 : find-permutation ( ... seq quot: ( ... elt -- ... ? ) -- ... elt/f )
 125     [ permutations-quot find drop ]
 126     [ drop over [ permutation ] [ 2drop f ] if ] 2bi ; inline
 127
 128 : reduce-permutations ( ... seq identity quot: ( ... prev elt -- ... next ) -- ... result )
 129     swapd each-permutation ; inline
 130
 131 : inverse-permutation ( seq -- permutation )
 132     <enumerated> sort-values keys ;
 133
 134 <PRIVATE
 135
 136 : cut-point ( seq -- n )
 137     [ last ] keep [ [ > ] keep swap ] find-last drop nip ; inline
 138
 139 : greater-from-last ( n seq -- i )
 140     [ nip ] [ nth ] 2bi [ > ] curry find-last drop ; inline
 141
 142 : reverse-tail! ( n seq -- seq )
 143     [ swap 1 + tail-slice reverse! drop ] keep ; inline
 144
 145 : (next-permutation) ( seq -- seq )
 146     dup cut-point [
 147         swap [ greater-from-last ] 2keep
 148         [ exchange ] [ reverse-tail! nip ] 3bi
 149     ] [ reverse! ] if* ;
 150
 151 HINTS: (next-permutation) array ;
 152
 153 PRIVATE>
 154
 155 : next-permutation ( seq -- seq )
 156     dup empty? [ (next-permutation) ] unless ;
 157
 158
 159 ! Combinadic-based combination methodology
 160
 161 <PRIVATE
 162
 163 ! "Algorithm 515: Generation of a Vector from the Lexicographical Index"
 164 ! Buckles, B. P., and Lybanon, M. ACM
 165 ! Transactions on Mathematical Software, Vol. 3, No. 2, June 1977.
 166
 167 :: combination-indices ( x! p n -- seq )
 168     x 1 + x!
 169     p 0 <array> :> c 0 :> k! 0 :> r!
 170     p 1 - [| i |
 171         i [ 0 ] [ 1 - c nth ] if-zero i c set-nth
 172         [ k x < ] [
 173             i c [ 1 + ] change-nth
 174             n i c nth - p i 1 + - nCk r!
 175             k r + k!
 176         ] do while k r - k!
 177     ] each-integer
 178     p 2 < [ 0 ] [ p 2 - c nth ] if
 179     p 1 < [ drop ] [ x + k - p 1 - c set-nth ] if
 180     c [ 1 - ] map! ;
 181
 182 PRIVATE>
 183
 184 : combination ( m seq k -- seq' )
 185     swap [ length combination-indices ] [ nths-unsafe ] bi ;
 186
 187 TUPLE: combinations seq k length ;
 188
 189 : <combinations> ( seq k -- combinations )
 190     2dup [ length ] [ nCk ] bi* combinations boa ;
 191
 192 M: combinations length length>> ; inline
 193 M: combinations nth-unsafe [ seq>> ] [ k>> ] bi combination ;
 194 M: combinations hashcode* tuple-hashcode ;
 195
 196 INSTANCE: combinations immutable-sequence
 197
 198 <PRIVATE
 199
 200 : find-max-index ( seq n -- i )
 201     over length - '[ _ + >= ] find-index drop ; inline
 202
 203 : increment-rest ( i seq -- )
 204     [ nth ] [ swap tail-slice ] 2bi
 205     [ drop 1 + dup ] map! 2drop ; inline
 206
 207 : increment-last ( seq -- )
 208     [ [ length 1 - ] keep [ 1 + ] change-nth ] unless-empty ; inline
 209
 210 :: next-combination ( seq n -- seq )
 211     seq n find-max-index [
 212         1 [-] seq increment-rest
 213     ] [
 214         seq increment-last
 215     ] if* seq ; inline
 216
 217 :: combinations-quot ( seq k quot -- seq quot' )
 218     seq length :> n
 219     n k nCk <iota> k <iota> >array seq quot n
 220     '[ drop _ [ _ nths-unsafe @ ] keep _ next-combination drop ] ; inline
 221
 222 PRIVATE>
 223
 224 : each-combination ( ... seq k quot: ( ... elt -- ... ) -- ... )
 225     combinations-quot each ; inline
 226
 227 : map-combinations ( ... seq k quot: ( ... elt -- ... newelt ) -- ... newseq )
 228     combinations-quot map ; inline
 229
 230 : filter-combinations ( ... seq k quot: ( ... elt -- ... ? ) -- ... newseq )
 231     selector [ each-combination ] dip ; inline
 232
 233 : map>assoc-combinations ( ... seq k quot: ( ... elt -- ... key value ) exemplar -- ... assoc )
 234     [ combinations-quot ] dip map>assoc ; inline
 235
 236 : all-combinations ( seq k -- seq' )
 237     [ ] map-combinations ;
 238
 239 : all-combinations? ( ... seq k quot: ( ... elt -- ... ? ) -- ... ? )
 240     combinations-quot all? ; inline
 241
 242 : find-combination ( ... seq k quot: ( ... elt -- ... ? ) -- ... elt/f )
 243     [ combinations-quot find drop ]
 244     [ drop pick [ combination ] [ 3drop f ] if ] 3bi ; inline
 245
 246 : reduce-combinations ( ... seq k identity quot: ( ... prev elt -- ... next ) -- ... result )
 247     [ -rot ] dip each-combination ; inline
 248
 249 : all-subsets ( seq -- subsets )
 250     dup length [0,b] [ all-combinations ] with map concat ;
 251
 252 <PRIVATE
 253
 254 :: next-selection ( seq n -- )
 255     1 seq length 1 - [
 256         dup 0 >= [ over 0 = ] [ t ] if
 257     ] [
 258         [ seq [ + n /mod ] change-nth-unsafe ] keep 1 -
 259     ] do until 2drop ; inline
 260
 261 :: (selections) ( seq n -- selections )
 262     seq length :> len
 263     n 0 <array> :> idx
 264     len n ^ [
 265         idx seq nths-unsafe
 266         idx len next-selection
 267     ] replicate ;
 268
 269 PRIVATE>
 270
 271 : selections ( seq n -- selections )
 272     dup 0 > [ (selections) ] [ 2drop { } ] if ;