add 'combinations' and 'in-combinations'

pkgs/racket-doc/scribblings/reference/pairs.scrbl

@@ -1252,9 +1252,34 @@ returns @racket[#f].
 Returns a list with all elements from @racket[lst], randomly shuffled.
 
 @mz-examples[#:eval list-eval
+  (shuffle '(1 2 3 4 5 6))
+  (shuffle '(1 2 3 4 5 6))
   (shuffle '(1 2 3 4 5 6))]}
 
 
+@defproc*[([(combinations [lst list?]) list?]
+           [(combinations [lst list?] [size exact-nonnegative-integer?]) list?])]{
+@margin-note{Wikipedia @hyperlink["https://en.wikipedia.org/wiki/Combination"]{combinations}}
+Return a list of all combinations of elements in the input list
+(aka the @index["powerset"]{powerset} of @racket[lst]).
+If @racket[size] is given, limit results to combinations of @racket[size] elements.
+
+@mz-examples[#:eval list-eval
+  (combinations '(1 2 3))
+  (combinations '(1 2 3) 2)]}
+
+
+@defproc*[([(in-combinations [lst list?]) list?]
+           [(in-combinations [lst list?] [size exact-nonnegative-integer?]) sequence?])]{
+@index["in-powerset"]{Returns} a sequence of all combinations of elements in the input list,
+ or all combinations of length @racket[size] if @racket[size] is given.
+Builds combinations one-by-one instead of all at once.
+
+@mz-examples[#:eval list-eval
+  (time (begin (combinations (range 15)) (void)))
+  (time (begin (in-combinations (range 15)) (void)))]}
+
+
 @defproc[(permutations [lst list?])
          list?]{
 

pkgs/racket-test-core/tests/racket/list.rktl

@@ -422,6 +422,42 @@
     (test expected length+sum (shuffle l)))
   (when (pair? l) (loop (cdr l))))
 
+;; ---------- combinations ----------
+(let ()
+  (define (comb<? l1 l2) ; (works only on tests with numeric lists)
+    (define L1 (length l1))
+    (define L2 (length l2))
+    (or (< L1 L2)
+        (and (= L1 L2)
+             (let loop ([l1 l1] [l2 l2])
+               (or (< (car l1) (car l2))
+                   (and (= (car l1) (car l2))
+                        (loop (cdr l1) (cdr l2))))))))
+  (define (sorted-combs l k)
+    (define l1 (sort (combinations l k) comb<?))
+    (define l2 (sort (for/list ([c (in-combinations l k)]) c) comb<?))
+    (test #t equal? l1 l2)
+    l1)
+  (test '(()) sorted-combs '() #f)
+  (test '(()) sorted-combs  '() 0)
+  (test '() sorted-combs '() 9)
+  (test '(() (6)) sorted-combs '(6) #f)
+  (test '(()) sorted-combs '(6) 0)
+  (test '((6)) sorted-combs '(6) 1)
+  (test '() sorted-combs '(6) 2)
+  (test '(() (8) (9) (9 8)) sorted-combs '(9 8) #f)
+  (test '((8) (9)) sorted-combs '(9 8) 1)
+  (test '((9 8)) sorted-combs '(9 8) 2)
+  (test
+   '(() (1) (2) (3) (4) (5) (1 2) (1 3) (1 5) (2 3) (2 5) (4 1) (4 2) (4 3) (4 5)
+     (5 3) (1 2 3) (1 2 5) (1 5 3) (2 5 3) (4 1 2) (4 1 3) (4 1 5) (4 2 3)
+     (4 2 5) (4 5 3) (1 2 5 3) (4 1 2 3) (4 1 2 5) (4 1 5 3) (4 2 5 3) (4 1 2 5 3))
+   sorted-combs '(4 1 2 5 3) #f)
+  (test '(()) sorted-combs '(4 1 2 5 3) 0)
+  (test
+   '((1 2) (1 3) (1 5) (2 3) (2 5) (4 1) (4 2) (4 3) (4 5) (5 3))
+   sorted-combs '(4 1 2 5 3) 2))
+
 ;; ---------- permutations ----------
 (let ()
   (define (perm<? l1 l2) ; (works only on tests with numeric lists)

racket/collects/racket/list.rkt

@@ -45,6 +45,8 @@
          append-map
          filter-not
          shuffle
+         combinations
+         in-combinations
          permutations
          in-permutations
          argmin
@@ -588,6 +590,84 @@
     (vector-set! a j x))
   (vector->list a))
 
+(define (combinations l [k #f])
+  (for/list ([x (in-combinations l k)]) x))
+
+;; Generate combinations of the list `l`.
+;; - If `k` is a natural number, generate all combinations of size `k`.
+;; - If `k` is #f, generate all combinations of any size (powerset of `l`).
+(define (in-combinations l [k #f])
+  (unless (list? l)
+    (raise-argument-error 'in-combinations "list?" 0 l))
+  (when (and k (not (exact-nonnegative-integer? k)))
+    (raise-argument-error 'in-combinations "exact-nonnegative-integer?" 1 k))
+  (define v (list->vector l))
+  (define N (vector-length v))
+  (define N-1 (- N 1))
+  (define gen-combinations
+    (cond
+      [(not k)
+       ;; Enumerate all binary numbers [1..2**N].
+       ;; Produce the combination with elements in `v` at the same
+       ;;  positions as the 1's in the binary number.
+       (define limit (expt 2 N))
+       (define curr-box (box 0))
+       (lambda ()
+         (let ([curr (unbox curr-box)])
+           (if (< curr limit)
+             (begin0
+               (for/fold ([acc '()])
+                         ([i (in-range N-1 -1 -1)])
+                 (if (bitwise-bit-set? curr i)
+                   (cons (vector-ref v i) acc)
+                   acc))
+               (set-box! curr-box (+ curr 1)))
+             #f)))]
+      [(< N k)
+       (lambda () #f)]
+      [else
+       ;; Keep a vector `k*` that contains `k` indices
+       ;; Use `k*` to generate combinations
+       (define k* #f) ; (U #f (Vectorof Index))
+       (define k-1 (- k 1))
+       ;; `k*-incr` tries to increment the positions in `k*`.
+       ;; On success, can use `k*` to build a combination.
+       ;; Returns #f on failure.
+       (define (k*-incr)
+         (cond
+          [(not k*)
+           ;; 1. Initialize the vector `k*` to the first {0..k-1} indices
+           (set! k* (build-vector k (lambda (i) i)))]
+          [(zero? k)
+           ;; (Cannot increment a zero vector)
+           #f]
+          [else
+           (or
+            ;; 2. Try incrementing the leftmost index that is
+            ;;    at least 2 less than the following index in `k*`.
+            (for/or ([i (in-range 0 k-1)])
+              (let ([k*_i   (vector-ref k* i)]
+                    [k*_i+1 (vector-ref k* (+ i 1))])
+                (and (< k*_i (- k*_i+1 1))
+                     (vector-set! k* i (+ k*_i 1)))))
+            ;; 3. Increment the rightmost index, up to a max of `N-1`.
+            ;;    Also replace the first `k-1` indices to `[0..k-2]`
+            (let ([k*_last (vector-ref k* k-1)])
+              (if (< k*_last N-1)
+                  (begin
+                    (vector-set! k* k-1 (+ k*_last 1))
+                    (for ([i (in-range k-1)])
+                      (vector-set! k* i i)))
+                  #f)))]))
+       (define (k*->combination)
+         ;; Get the `k` elements indexed by `k*`
+         (for/fold ([acc '()])
+                   ([i (in-range k-1 -1 -1)])
+           (cons (vector-ref v (vector-ref k* i)) acc)))
+       (lambda ()
+         (and (k*-incr) (k*->combination)))]))
+      (in-producer gen-combinations #f))
+
 ;; This implements an algorithm known as "Ord-Smith".  (It is described in a
 ;; paper called "Permutation Generation Methods" by Robert Sedgewlck, listed as
 ;; Algorithm 8.)  It has a number of good properties: it is very fast, returns