fix (combinations n k) bug

Now using Gosper's hack to enumerate length k binary numbers.
New implementation is shorter & a little more obviously correct
(if you trust the bit-twiddling)
https://en.wikipedia.org/wiki/Combinatorial_number_system#Applications
This commit is contained in:
ben 2016-02-16 01:03:09 -05:00
parent 301b47df2c
commit f83cec1b04
2 changed files with 43 additions and 63 deletions

View File

@ -456,7 +456,16 @@
(test '(()) sorted-combs '(4 1 2 5 3) 0) (test '(()) sorted-combs '(4 1 2 5 3) 0)
(test (test
'((1 2) (1 3) (1 5) (2 3) (2 5) (4 1) (4 2) (4 3) (4 5) (5 3)) '((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)) sorted-combs '(4 1 2 5 3) 2)
(test
'((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))
sorted-combs '(4 1 2 5 3) 3)
(test
21
(lambda (n k)
(length (combinations n k)))
'(1 2 3 4 5 6 7)
5))
;; ---------- permutations ---------- ;; ---------- permutations ----------
(let () (let ()

View File

@ -604,69 +604,40 @@
(define v (list->vector l)) (define v (list->vector l))
(define N (vector-length v)) (define N (vector-length v))
(define N-1 (- N 1)) (define N-1 (- N 1))
(define gen-combinations (define (vector-ref/bits v b)
(for/fold ([acc '()])
([i (in-range N-1 -1 -1)])
(if (bitwise-bit-set? b i)
(cons (vector-ref v i) acc)
acc)))
(define-values (first last incr)
(cond (cond
[(not k) [(not k)
;; Enumerate all binary numbers [1..2**N]. ;; Enumerate all binary numbers [1..2**N].
;; Produce the combination with elements in `v` at the same (values 0 (- (expt 2 N) 1) add1)]
;; positions as the 1's in the binary number. [(< N k)
(define limit (expt 2 N)) ;; Nothing to produce
(define curr-box (box 0)) (values 1 0 values)]
[else
;; Enumerate numbers with `k` ones, smallest to largest
(define first (- (expt 2 k) 1))
(define gospers-hack ;; https://en.wikipedia.org/wiki/Combinatorial_number_system#Applications
(if (zero? first)
add1
(lambda (n)
(let* ([u (bitwise-and n (- n))]
[v (+ u n)])
(+ v (arithmetic-shift (quotient (bitwise-xor v n) u) -2))))))
(values first (arithmetic-shift first (- N k)) gospers-hack)]))
(define gen-next
(let ([curr-box (box first)])
(lambda () (lambda ()
(let ([curr (unbox curr-box)]) (let ([curr (unbox curr-box)])
(if (< curr limit) (and (<= curr last)
(begin0 (begin0
(for/fold ([acc '()]) (vector-ref/bits v curr)
([i (in-range N-1 -1 -1)]) (set-box! curr-box (incr curr))))))))
(if (bitwise-bit-set? curr i) (in-producer gen-next #f))
(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 ;; This implements an algorithm known as "Ord-Smith". (It is described in a
;; paper called "Permutation Generation Methods" by Robert Sedgewlck, listed as ;; paper called "Permutation Generation Methods" by Robert Sedgewlck, listed as