advent of code (2018) day 23

day23.rkt

#lang racket

(require data/heap)

(define in (open-input-file "day23.in"))
(define line-pattern #px"pos=<([-[:digit:]]+),([-[:digit:]]+),([-[:digit:]]+)>, r=([-[:digit:]]+)")

(define bots
  (for/list ([line (in-lines in)])
    (let ([match (map string->number
                      (rest (regexp-match line-pattern line)))])
      (cons (take match 3) (last match)))))

(struct octree (min-x max-x min-y max-y min-z max-z))

(define (octree-expand tree bot)
  (match-let ([(struct octree (min-x max-x min-y max-y min-z max-z)) tree]
              [(cons (list x y z) r) bot])
    (octree
     (min min-x x) (max max-x x)
     (min min-y y) (max max-y y)
     (min min-z z) (max max-z z))))

(define (octree-leaf? tree)
  (match-let ([(struct octree (min-x max-x min-y max-y min-z max-z)) tree])
    (and (= min-x max-x) (= min-y max-y) (= min-z max-z))))

(define (octree-contains? tree bot)
  (match-let ([(struct octree (min-x max-x min-y max-y min-z max-z)) tree]
              [(cons (list x y z) r) bot])
    ; furthest distance from the bot to the tree
    (let ([furthest-distance
           (+ (max (abs (- x min-x)) (abs (- x max-x)))
              (max (abs (- y min-y)) (abs (- y max-y)))
              (max (abs (- z min-z)) (abs (- z max-z))))])
      (>= r furthest-distance))))

(define (octree-disjoint? tree bot)
  (define (interval-min left right point)
    (or (and (<= left point) (>= right point) 0)
        (min (abs (- point left)) (abs (- point right)))))
  (match-let ([(struct octree (min-x max-x min-y max-y min-z max-z)) tree]
              [(cons (list x y z) r) bot])
    (let ([closest-distance
           (+ (interval-min min-x max-x x)
              (interval-min min-y max-y y)
              (interval-min min-z max-z z))])
      (< r closest-distance))))

(define (octree-num-bots tree bots)
  (count (lambda (bot)
           (or (octree-contains? tree bot)
               (not (octree-disjoint? tree bot))))
         bots))

(define (octree-distance tree)
  (unless (octree-leaf? tree) (error "called octree-distance on a non-leaf tree"))
  (+ (abs (octree-min-x tree)) (abs (octree-min-y tree)) (abs (octree-min-z tree))))

(define (octree-subdivide tree bots)
  (define (mid a b) (quotient (+ a b 1) 2))
  (match-let ([(struct octree (min-x max-x min-y max-y min-z max-z)) tree])
    (let ([mid-x (mid min-x max-x)]
          [mid-y (mid min-y max-y)]
          [mid-z (mid min-z max-z)])
      (for*/list ([(min-x max-x) (in-parallel (list min-x mid-x) (list (sub1 mid-x) max-x))]
                  #:when (>= max-x min-x)
                  [(min-y max-y) (in-parallel (list min-y mid-y) (list (sub1 mid-y) max-y))]
                  #:when (>= max-y min-y)
                  [(min-z max-z) (in-parallel (list min-z mid-z) (list (sub1 mid-z) max-z))]
                  #:when (>= max-z min-z)
                  [tree (in-value (octree min-x max-x min-y max-y min-z max-z))])
        (cons tree (octree-num-bots tree bots))))))
         

(define (best-point bots)
  ; binary heap sorted by number of bots in range of
  ; octree nodes
  ; elements are (octree . num-bots) pairs
  (define heap (make-heap (lambda (a b) (> (cdr a) (cdr b)))))
  ; loop with the current record as a (num-bots . origin-distance) pair
  (define (loop cur-best iter)
    (cond
      [(zero? (heap-count heap)) cur-best] ; done
      [else
       ; top element is a (non-leaf) octree node
       (match-let ([(cons tree num-bots) (heap-min heap)])
         (heap-remove-min! heap)
         (cond
           ; cull the node if it's not better than the current best
           [(and cur-best (< num-bots (car cur-best)))
            (loop cur-best (add1 iter))]
           ; otherwise, we're at a node we can subdivide
           [else
            (let* ([new-trees (octree-subdivide tree bots)]
                   [new-best
                    (for/fold ([new-best cur-best])
                              ([t new-trees]
                               #:when (or (not new-best) (>= (cdr t) (car new-best))))
                      (cond
                        [(octree-leaf? (car t))
                         ; leaf node, compare to best
                         (if (or (not new-best)
                                 (> (cdr t) (car new-best))
                                 (< (octree-distance (car t)) (cdr new-best)))
                             (cons (cdr t) (octree-distance (car t)))
                             new-best)]
                        [else
                         ; non-leaf node, add to the heap
                         ; we already pruned it against new-best
                         (heap-add! heap t)
                         new-best]))])
              (loop new-best (add1 iter)))]))]))
  (let ([initial-tree
         (for/fold ([tree (octree 0 0 0 0 0 0)])
                   ([bot bots])
           (octree-expand tree bot))])
    (heap-add! heap (cons initial-tree (length bots)))
    (cdr (loop #f 0))))