wow

day6.rkt

@@ -1,43 +1,34 @@
 #lang racket
 (require "lib/common.rkt"
-         fancy-app)
-
-(define (generate-initial-hash fish)
-  (define res (make-hash))
+         fancy-app
+         math/array
+         math/matrix)
+
+(define transition-matrix
+  (matrix [[0 1 0 0 0 0 0 0 0]
+           [0 0 1 0 0 0 0 0 0]
+           [0 0 0 1 0 0 0 0 0]
+           [0 0 0 0 1 0 0 0 0]
+           [0 0 0 0 0 1 0 0 0]
+           [0 0 0 0 0 0 1 0 0]
+           [1 0 0 0 0 0 0 1 0]
+           [0 0 0 0 0 0 0 0 1]
+           [1 0 0 0 0 0 0 0 0]]))
+
+(define (generate-initial-vector fish)
+  (define vec (make-vector 9 0))
   (for ([f (in-list fish)])
-    (hash-update! res f add1 0))
-  res)
-
-(define (simulate-day! hsh)
-  (define copy (hash-copy hsh))
-  (for ([(f k) (in-hash copy)])
-    (cond [(zero? f)
-           ;; add k 8s
-           (hash-update! hsh 8 (+ k _) 0)
-           ;; add k 6s
-           (hash-update! hsh 6 (+ k _) 0)
-           ;; reset zero
-           (hash-set! hsh 0 0)]
-          [else
-           ;; move it down
-           (hash-update! hsh (- f 1) (+ k _) 0)
-           (hash-update! hsh f (- _ k) 0)])))
-
-(define (day6a fish)
-  (define count-hash (generate-initial-hash fish))
-
-  (for ([_ (in-range 80)])
-    (simulate-day! count-hash))
-
-  (foldl + 0 (hash-values count-hash)))
-
-(define (day6b fish)
-  (define count-hash (generate-initial-hash fish))
+    (vector-set! vec f (add1 (vector-ref vec f))))
+  vec)
 
-  (for ([_ (in-range 256)])
-    (simulate-day! count-hash))
+(define (simulate fish days)
+  (define V
+    (matrix* (matrix-expt transition-matrix days)
+             (->col-matrix (generate-initial-vector fish))))
+  (for/sum ([v (in-array V)]) v))
 
-  (foldl + 0 (hash-values count-hash)))
+(define (day6a fish) (simulate fish 80))
+(define (day6b fish) (simulate fish 256))
 
 (module+ main
   (call-with-input-file "data/day6.txt"
@@ -54,45 +45,6 @@
 (module+ test
   (require rackunit)
   (define test-fish '(3 4 3 1 2))
-  (define init-hash (generate-initial-hash test-fish))
-  (check-equal? (hash->list init-hash)
-                (hash->list (hash 1 1
-                                  2 1
-                                  3 2
-                                  4 1)))
-  (simulate-day! init-hash)
-  (check-equal? (sort #:key car (hash->list init-hash) <)
-                (sort #:key car
-                      (hash->list (hash 0 1
-                                        1 1
-                                        2 2
-                                        3 1
-                                        4 0))
-                      <))
-  (simulate-day! init-hash)
-  (check-equal? (sort #:key car (hash->list init-hash) <)
-                (sort #:key car
-                      (hash->list (hash 0 1
-                                        1 2
-                                        2 1
-                                        3 0
-                                        4 0
-                                        6 1
-                                        8 1))
-                      <))
-  (simulate-day! init-hash)
-  (check-equal? (sort #:key car (hash->list init-hash) <)
-                (sort #:key car
-                      (hash->list (hash 0 2
-                                        1 1
-                                        2 0
-                                        3 0
-                                        4 0
-                                        5 1
-                                        6 1
-                                        7 1
-                                        8 1))
-                      <))
 
   (check-equal? (day6a test-fish) 5934)
   (check-equal? (day6b test-fish) 26984457539))