From 2ad2b81998ab73365ca37d960195d12049ddcb1f Mon Sep 17 00:00:00 2001
From: John Benediktsson <mrjbq7@gmail.com>
Date: Fri, 28 Jan 2022 09:23:22 -0800
Subject: [PATCH] math.combinatorics: adding combinations-with-replacement

---
 .../combinatorics/combinatorics-tests.factor  |  2 +
 basis/math/combinatorics/combinatorics.factor | 39 +++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/basis/math/combinatorics/combinatorics-tests.factor b/basis/math/combinatorics/combinatorics-tests.factor
index e18771e333..3111d51bbd 100644
--- a/basis/math/combinatorics/combinatorics-tests.factor
+++ b/basis/math/combinatorics/combinatorics-tests.factor
@@ -52,6 +52,8 @@ math.combinatorics.private tools.test sequences sets ;
 { "ABC" "ACB" "BAC" "BCA" "CAB" "CBA" "ABC" }
 [ "ABC" 6 [ dup dup clone-like next-permutation ] times ] unit-test
 
+{ { "AA" "AB" "AC" "BB" "BC" "CC" } } [ "ABC" 2 all-combinations-with-replacement ] unit-test
+
 { { 0 1 2 } } [ 0 3 5 combination-indices ] unit-test
 { { 2 3 4 } } [ 9 3 5 combination-indices ] unit-test
 
diff --git a/basis/math/combinatorics/combinatorics.factor b/basis/math/combinatorics/combinatorics.factor
index f4e9d5f2d3..d4023fa82a 100644
--- a/basis/math/combinatorics/combinatorics.factor
+++ b/basis/math/combinatorics/combinatorics.factor
@@ -184,6 +184,45 @@ PRIVATE>
 
 <PRIVATE
 
+:: next-combination-with-replacement ( seq n -- seq )
+    seq n 1 - '[ _ = not ] find-last drop :> i
+    seq i tail-slice i seq nth 1 + '[ drop _ ] map! drop
+    seq ; inline
+
+:: combinations-with-replacement-quot ( seq k quot -- seq quot' )
+    seq length :> n
+    k 1 - n + factorial k factorial / n 1 - factorial / <iota>
+    k 0 <array> seq quot n
+    '[ drop _ [ _ nths-unsafe @ ] keep _ next-combination-with-replacement drop ] ; inline
+
+PRIVATE>
+
+: each-combination-with-replacement ( ... seq k quot: ( ... elt -- ... ) -- ... )
+    combinations-with-replacement-quot each ; inline
+
+: map-combinations-with-replacement ( ... seq k quot: ( ... elt -- ... newelt ) -- ... newseq )
+    combinations-with-replacement-quot map ; inline
+
+: filter-combinations-with-replacement ( ... seq k quot: ( ... elt -- ... ? ) -- ... newseq )
+    selector [ each-combination-with-replacement ] dip ; inline
+
+: map>assoc-combinations-with-replacement ( ... seq k quot: ( ... elt -- ... key value ) exemplar -- ... assoc )
+    [ combinations-with-replacement-quot ] dip map>assoc ; inline
+
+: all-combinations-with-replacement ( seq k -- seq' )
+    [ ] map-combinations-with-replacement ;
+
+: all-combinations-with-replacement? ( ... seq k quot: ( ... elt -- ... ? ) -- ... ? )
+    combinations-with-replacement-quot all? ; inline
+
+: find-combination-with-replacement ( ... seq k quot: ( ... elt -- ... ? ) -- ... elt/f )
+    [ f ] 3dip '[ nip _ keep swap ] combinations-with-replacement-quot find drop swap and ; inline
+
+: reduce-combinations-with-replacement ( ... seq k identity quot: ( ... prev elt -- ... next ) -- ... result )
+    -rotd each-combination-with-replacement ; inline
+
+<PRIVATE
+
 ! "Algorithm 515: Generation of a Vector from the Lexicographical Index"
 ! Buckles, B. P., and Lybanon, M. ACM
 ! Transactions on Mathematical Software, Vol. 3, No. 2, June 1977.
-- 
2.34.1