diff --git a/pkgs/racket-doc/scribblings/reference/pairs.scrbl b/pkgs/racket-doc/scribblings/reference/pairs.scrbl index fa75e3b11e..fdc6e24f5a 100644 --- a/pkgs/racket-doc/scribblings/reference/pairs.scrbl +++ b/pkgs/racket-doc/scribblings/reference/pairs.scrbl @@ -1221,6 +1221,33 @@ result of @racket[proc]. Signals an error on an empty list. (argmax car '((3 pears) (3 oranges)))]} +@defproc[(group-by [key (-> any/c any/c)] + [lst list?] + [same? (any/c any/c . -> . any/c) equal?]) + (listof list?)]{ + +Groups the given list into equivalence classes, with equivalence being +determined by @racket[same?]. Within each equivalence class, @racket[group-by] +preserves the ordering of the original list. Equivalence classes themselves are +in order of first appearance in the input. + +@examples[#:eval list-eval +(group-by (lambda (x) (modulo x 3)) '(1 2 1 2 54 2 5 43 7 2 643 1 2 0)) +] +} + +@defproc[(cartesian-product [lst list?] ...) + (listof list?)]{ + +Computes the n-ary cartesian product of the given lists. + +@examples[#:eval list-eval +(cartesian-product '(1 2 3) '(a b c)) +(cartesian-product '(4 5 6) '(d e f) '(#t #f)) +] +} + + @close-eval[list-eval] diff --git a/pkgs/racket-test-core/tests/racket/list.rktl b/pkgs/racket-test-core/tests/racket/list.rktl index 55d1ca3e41..539328802e 100644 --- a/pkgs/racket-test-core/tests/racket/list.rktl +++ b/pkgs/racket-test-core/tests/racket/list.rktl @@ -540,4 +540,31 @@ (test '(20 19 18 17 16 15 14 13 12 11) range 20 10 -1) (test '(10 11.5 13.0 14.5) range 10 15 1.5)) +;; ---------- group-by ---------- + +(test '((1) (4) (2 2) (56) (3)) group-by values '(1 4 2 56 2 3)) +(test '((1 1 1) (2 2 2 2 2) (54) (5) (43) (7) (643) (0)) + group-by values '(1 2 1 2 54 2 5 43 7 2 643 1 2 0)) +(test '((1 3) (4 2 56 2)) + group-by values '(1 4 2 56 2 3) (lambda (x y) (or (and (even? x) (even? y)) + (and (odd? x) (odd? y))))) +(test '(((1 a)) ((4 b)) ((2 c) (2 e)) ((56 d)) ((3 f))) + group-by car '((1 a) (4 b) (2 c) (56 d) (2 e) (3 f))) +(test '((1 3 5) (2 4 6)) group-by even? '(1 2 3 4 5 6)) +(err/rt-test (group-by #f)) +(err/rt-test (group-by '() #f)) +(err/rt-test (group-by '() values #f)) + +;; ---------- cartesian-product ---------- + +(test '((1 a) (1 b) (1 c) + (2 a) (2 b) (2 c) + (3 a) (3 b) (3 c)) + cartesian-product '(1 2 3) '(a b c)) +(test '((4 d #t) (4 d #f) (4 e #t) (4 e #f) (4 f #t) (4 f #f) + (5 d #t) (5 d #f) (5 e #t) (5 e #f) (5 f #t) (5 f #f) + (6 d #t) (6 d #f) (6 e #t) (6 e #f) (6 f #t) (6 f #f)) + cartesian-product '(4 5 6) '(d e f) '(#t #f)) +(err/rt-test (cartesian-product 3)) + (report-errs) diff --git a/racket/collects/racket/list.rkt b/racket/collects/racket/list.rkt index b0ffde536b..838d64a9a3 100644 --- a/racket/collects/racket/list.rkt +++ b/racket/collects/racket/list.rkt @@ -40,7 +40,9 @@ permutations in-permutations argmin - argmax) + argmax + group-by + cartesian-product) (define (first x) (if (and (pair? x) (list? x)) @@ -611,3 +613,65 @@ (loop min min-var (cdr xs))]))])))) (define (argmin f xs) (mk-min < 'argmin f xs)) (define (argmax f xs) (mk-min > 'argmax f xs)) + +;; (x -> y) (listof x) [(y y -> bool)] -> (listof (listof x)) +;; groups together elements that are considered equal +;; =? should be reflexive, transitive and commutative +(define (group-by key l [=? equal?]) + + (unless (and (procedure? key) + (procedure-arity-includes? key 1)) + (raise-argument-error 'group-by "(-> any/c any/c)" key)) + (unless (and (procedure? =?) + (procedure-arity-includes? =? 2)) + (raise-argument-error 'group-by "(any/c any/c . -> . any/c)" =?)) + (unless (list? l) + (raise-argument-error 'group-by "list?" l)) + + ;; like hash-update, but for alists + (define (alist-update al k up fail) + (let loop ([al al]) + (cond [(null? al) + ;; did not find equivalence class, create one + (list (cons k (up '())))] + [(=? (car (car al)) k) + ;; found the right equivalence class + (cons + (cons k (up (cdr (car al)))) ; updater takes elements, w/o key + (cdr al))] + [else ; keep going + (cons (car al) (loop (cdr al)))]))) + + ;; In cases where `=?` is a built-in equality, can use hash tables instead + ;; of lists to compute equivalence classes. + (define-values (base update) + (cond [(equal? =? eq?) (values (hasheq) hash-update)] + [(equal? =? eqv?) (values (hasheqv) hash-update)] + [(equal? =? equal?) (values (hash) hash-update)] + [else (values '() alist-update)])) + + (define classes + (for/fold ([res base]) + ([elt (in-list l)] + [idx (in-naturals)]) ; to keep ordering stable + (define k (key elt)) + (define v (cons idx elt)) + (update res k (lambda (o) (cons v o)) '()))) + (define sorted-classes + (if (list? classes) + (for/list ([p (in-list classes)]) + (sort (cdr p) < #:key car)) + (for/list ([(_ c) (in-hash classes)]) + (sort c < #:key car)))) + ;; sort classes by order of first appearance, then remove indices + (for/list ([c (in-list (sort sorted-classes < #:key caar))]) + (map cdr c))) + +;; (listof x) ... -> (listof (listof x)) +(define (cartesian-product . ls) + (for ([l (in-list ls)]) + (unless (list? l) + (raise-argument-error 'cartesian-product "list?" l))) + (define (cp-2 as bs) + (for*/list ([i (in-list as)] [j (in-list bs)]) (cons i j))) + (foldr cp-2 (list (list)) ls))