aoc2020 day7

aoc2020-day7.lisp

;;;; +----------------------------------------------------------------+
;;;; | Advent of Code 2020                                            |
;;;; +----------------------------------------------------------------+

(defpackage #:snippets/aoc2020/day7
  (:use #:cl)
  (:import-from
   #:fset
   #:empty-map
   #:with
   #:lookup
   #:range
   #:do-set)
  (:import-from
   #:split-sequence
   #:split-sequence-if)
  (:import-from
   #:alexandria
   #:make-keyword)
  (:import-from
   #:org.tfeb.hax.memoize
   #:def-memoized-function
   #:clear-memoized-function)
  (:export
   #:day7))

(in-package #:snippets/aoc2020/day7)

;;;; Bag

(defstruct (bag (:print-function print-bag))
  color
  contained-by
  contains)

(defun print-bag (bag stream depth)
  (declare (ignore depth))
  (format stream "#<BAG ~S>" (bag-color bag)))

;;;; Parser

;;; The following grammar is implemented:
;;;
;;; bag-spec ::= color-spec "contain" amount-list
;;;
;;; amount-list ::= "no" "other" "bags"
;;;               | amount-spec*
;;;
;;; amount-spec ::= number color-spec
;;;
;;; color-spec ::= color "bags"
;;;              | color "bag"

(defun parse (input)
  (let ((rules (reduce #'note-bag-spec input :initial-value (empty-map))))
    (do-set (bag (range rules))
      (setf (bag-contains bag)
            (loop for (n color) in (bag-contains bag)
                  collect (list n
                                (or (lookup rules color)
                                    (error "Can't find bag of color ~S." color)))))
      (loop for (n child) in (bag-contains bag)
            do (push bag (bag-contained-by child))))
    rules))

(defun note-bag-spec (map spec)
  (multiple-value-bind (color contains)
      (parse-bag-spec (bag-spec-tokens spec))
    (with map color (make-bag :color color :contains contains))))

(defun bag-spec-tokens (spec)
  (split-sequence-if (lambda (char)
                       (find char " ,."))
                     spec
                     :remove-empty-subseqs t))

(defun parse-bag-spec (tokens)
  (multiple-value-bind (color more)
      (parse-color-spec tokens)
    (assert (equal (first more) "contain"))
    (values color
            (parse-amount-list (rest more)))))

(defun parse-color-spec (tokens)
  (let ((more (member-if (lambda (token)
                           (or (equal "bag" token)
                               (equal "bags" token)))
                         tokens)))
    (values (make-color (ldiff tokens more))
            (rest more))))

(defun make-color (tokens)
  (make-keyword (format nil "~@:(~{~A~^-~}~)" tokens)))

(defun parse-amount-list (tokens)
  (if (equal tokens '("no" "other" "bags"))
      (values nil nil)
      (parse-amount-spec* tokens)))

(defun parse-amount-spec* (tokens)
  (when tokens
    (multiple-value-bind (amount more)
        (parse-amount-spec tokens)
      (cons amount (parse-amount-spec* more)))))

(defun parse-amount-spec (tokens)
  (let ((number (parse-integer (first tokens))))
    (multiple-value-bind (color more)
        (parse-color-spec (rest tokens))
      (values (list number color)
              more))))

;;;; Queries

(def-memoized-function ancestors (bag)
  (reduce (lambda (set parent)
            (union set (ancestors parent)))
          (bag-contained-by bag)
          :initial-value (bag-contained-by bag)))

(defun num-ancestors (bag)
  (length (ancestors bag)))

(def-memoized-function total-contained (bag)
  (loop for (n child) in (bag-contains bag)
        sum (* n (1+ (total-contained child)))))

(defun day7 (input)
  (let ((rules (parse input)))
    (mapc #'clear-memoized-function '(ancestors total-contained))
    (list (num-ancestors (lookup rules :shiny-gold))
          (total-contained (lookup rules :shiny-gold)))))