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