diff --git a/pkgs/racket-doc/scribblings/reference/pairs.scrbl b/pkgs/racket-doc/scribblings/reference/pairs.scrbl index a39319377d..bad5ca7aa5 100644 --- a/pkgs/racket-doc/scribblings/reference/pairs.scrbl +++ b/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?]{ diff --git a/pkgs/racket-test-core/tests/racket/list.rktl b/pkgs/racket-test-core/tests/racket/list.rktl index cc185354c4..51eb4b4452 100644 --- a/pkgs/racket-test-core/tests/racket/list.rktl +++ b/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 (comblist 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