day 22: some cleanup/optimization

22.rkt

@@ -3,63 +3,51 @@
 (require "scripts/aoc.rkt")
 
 ;; solution for day 22
-  (define (score cards)
+(define (score cards)
-    (for/sum ([card (in-list cards)] [mult (in-range (length cards) 0 -1)])
+  (for/sum ([card (in-list cards)] [mult (in-range (length cards) 0 -1)])
-      (* card mult)))
+    (* card mult)))
 
-  (define (play p1 p2)
+(define (play p1 p2)
-    (cond
+  (match* (p1 p2)
-      [(empty? p1) (score p2)]
+    [('() _) (score p2)]
-      [(empty? p2) (score p1)]
+    [(_ '()) (score p1)]
-      [else
+    [((cons p1-card p1-rest) (cons p2-card p2-rest))
-       (define p1-card (first p1))
+     (if (> p1-card p2-card)
-       (define p2-card (first p2))
+       (play (append p1-rest (list p1-card p2-card)) p2-rest)
-       (cond
+       (play p1-rest (append p2-rest (list p2-card p1-card))))]))
-        [(> p1-card p2-card) (play (append (rest p1) (list p1-card p2-card)) (rest p2))]
-        [else (play (rest p1) (append (rest p2) (list p2-card p1-card)))])]))
 
 (define (part1 input)
   (match-define (list p1 p2) input)
-
   (play p1 p2))
 
-
 (define (part2 input)
   (match-define (list p1 p2) input)
 
-  (define memo (make-hash))
+  (define/memoized (play/recur p1 p2)
-  (define (do-memo level p1 p2)
+    (define local-memo (mutable-set))
-    (hash-update! memo level (lambda (x) (set-add x (list p1 p2))) (set)))
+    (let loop ([p1 p1] [p2 p2])
-
+      (cond
-  (define (play/recur p1 p2 [level 0])
+        [(set-member? local-memo (list p1 p2)) (cons 'p1 p1)]
-    ; (printf "~a ~a\n" p1 p2)
+        [(empty? p1) (cons 'p2 p2)]
-    (cond
+        [(empty? p2) (cons 'p1 p1)]
-      [(set-member? (hash-ref memo level (set)) (list p1 p2)) (cons 'p1 (score p1))]
-      [(empty? p1) (cons 'p2 (score p2))]
-      [(empty? p2) (cons 'p1 (score p1))]
-      [else
-       (define p1-card (first p1))
-       (define p1-rest (rest p1))
-       (define p2-card (first p2))
-       (define p2-rest (rest p2))
-       (cond
-        [(and (>= (length p1-rest) p1-card)
-              (>= (length p2-rest) p2-card))
-         (match-define (cons winner score)
-           (play/recur (take p1-rest p1-card) (take p2-rest p2-card) (add1 level)))
-         ; (printf "submatch ~a\n" winner)
-         (do-memo level p1 p2)
-         (if (symbol=? winner 'p1)
-           (play/recur (append p1-rest (list p1-card p2-card)) p2-rest level)
-           (play/recur p1-rest (append p2-rest (list p2-card p1-card)) level))]
-        [(> p1-card p2-card)
-         (do-memo level p1 p2)
-         (play/recur (append p1-rest (list p1-card p2-card)) p2-rest level)]
         [else
-          (do-memo level p1 p2)
+         (set-add! local-memo (list p1 p2))
-           (play/recur p1-rest (append p2-rest (list p2-card p1-card)) level)])]))
+         (match-define (cons p1-card p1-rest) p1)
+         (match-define (cons p2-card p2-rest) p2)
+         (cond
+          [(and (>= (length p1-rest) p1-card)
+                (>= (length p2-rest) p2-card))
+           (match-define (cons winner score)
+             (play/recur (take p1-rest p1-card) (take p2-rest p2-card)))
+           (if (symbol=? winner 'p1)
+             (loop (append p1-rest (list p1-card p2-card)) p2-rest)
+             (loop p1-rest (append p2-rest (list p2-card p1-card))))]
+          [(> p1-card p2-card)
+           (loop (append p1-rest (list p1-card p2-card)) p2-rest)]
+          [else
+           (loop p1-rest (append p2-rest (list p2-card p1-card)))])])))
 
-  (cdr (play/recur p1 p2)))
+  (score (cdr (play/recur p1 p2))))
 
 ;; parse input file
 (define (parse fname)
@@ -68,6 +56,8 @@
     (map string->number (rest (string-split p1 "\n")))
     (map string->number (rest (string-split p2 "\n")))))
 
+(time (part2 (parse "/tmp/adventofcode-2020/input/22.txt")))
+
 (module+ test
   (require rackunit)
   ;; tests here

scripts/aoc.rkt

@@ -26,11 +26,11 @@
 ;; makes a memoization wrapper around a function
 (define (memoize func)
   (define memo (make-hash))
-  (lambda (arg)
+  (lambda args
-    (hash-ref! memo arg (λ () (func arg)))))
+    (hash-ref! memo args (λ () (apply func args)))))
 
-(define-simple-macro (define/memoized (name:id arg:id) body:expr ...+)
+(define-simple-macro (define/memoized (name:id arg:id ...) body:expr ...+)
-  (define name (memoize (λ (arg) body ...))))
+  (define name (memoize (λ (arg ...) body ...))))
 
 ;; helper for maximum-bipartite-matching
 ;; given an edge list where each first item corresponds to the left set and each second item