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)]
(lambda () [else
(let ([curr (unbox curr-box)]) ;; Enumerate numbers with `k` ones, smallest to largest
(if (< curr limit) (define first (- (expt 2 k) 1))
(begin0 (define gospers-hack ;; https://en.wikipedia.org/wiki/Combinatorial_number_system#Applications
(for/fold ([acc '()]) (if (zero? first)
([i (in-range N-1 -1 -1)]) add1
(if (bitwise-bit-set? curr i) (lambda (n)
(cons (vector-ref v i) acc) (let* ([u (bitwise-and n (- n))]
acc)) [v (+ u n)])
(set-box! curr-box (+ curr 1))) (+ v (arithmetic-shift (quotient (bitwise-xor v n) u) -2))))))
#f)))] (values first (arithmetic-shift first (- N k)) gospers-hack)]))
[(< N k) (define gen-next
(lambda () #f)] (let ([curr-box (box first)])
[else (lambda ()
;; Keep a vector `k*` that contains `k` indices (let ([curr (unbox curr-box)])
;; Use `k*` to generate combinations (and (<= curr last)
(define k* #f) ; (U #f (Vectorof Index)) (begin0
(define k-1 (- k 1)) (vector-ref/bits v curr)
;; `k*-incr` tries to increment the positions in `k*`. (set-box! curr-box (incr curr))))))))
;; On success, can use `k*` to build a combination. (in-producer gen-next #f))
;; 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