9.rkt

#lang racket
(require threading)

(define FILE "/tmp/input")

(struct grid (data length width) #:transparent)

(define/contract (point? x)
  (-> any/c boolean?)
  (match x
    [(list (? number?) (? number?)) #t]
    [_ #f]))

(define (make-grid vecs)
  (grid
   vecs
   (vector-length vecs)
   (vector-length (vector-ref vecs 0))))

(define (char->number c)
  (string->number (string c)))

(define (parse-line line)
  (~> line
      string-trim
      string->list
      (map char->number _)
      (apply vector _)))

(define board
  (~>
   (with-input-from-file
       FILE
     (λ ()
       (port->lines (current-input-port))))
   (map parse-line _)
   (apply vector _)
   make-grid))

(define/contract (access p)
  (-> point? (or/c number? false?))
  (if (in-bounds? p)     
      (~> board
          grid-data
          (vector-ref (first p))
          (vector-ref (second p)))
      #f))

(define/contract (in-bounds? p)
  (-> point? boolean?)
  (define x (first p))
  (define y (second p))
  (and
   (>= x 0)
   (>= y 0)
   (< x (grid-length board))
   (< y (grid-width board))))

(define/contract (add-coord p1)
  (-> point? (-> point? point?))
  (λ (p2)
    (list (+ (first p1) (first p2))
          (+ (second p1) (second p2)))))
             

(define/contract (get-neighbors p)
  (-> point? (listof point?))
  (define transforms
    '((1 0)
      (0 1)
      (-1 0)
      (0 -1)))
  (~> transforms
      (map (add-coord p) _)
      (filter in-bounds? _)))

(define/contract (point<? p1)
  (-> point? (-> point? boolean?))
  (λ (p2)
    (< (access p1) (access p2))))

(define/contract (point>? p1)
  (-> point? (-> point? boolean?))
  (λ (p2)
    ((access p1) . < . (access p2))))

(define/contract (low-point? p)
  (-> point? (or/c point? false?))
  (define x (access p))
  (define neighbors (map access (get-neighbors p)))
  (if (andmap (λ (y) (< x y)) neighbors)
      p
      #f))
                    
(define/contract (all-cords)
  (-> (listof point?))
  (cartesian-product
   (stream->list (in-range (grid-length board)))
   (stream->list (in-range (grid-width board)))))

(define (in-basin? cur)
  (λ (next)
    (and (not (= (access next) 9))
         (< (access cur) (access next)))))

(define/contract (expand-basin p)
  (-> point? (listof point?))
  (define ns (filter (in-basin? p) (get-neighbors p)))
  (append (cons p '())
          (apply append
                 (for/list [(n (in-list ns))]
                   (expand-basin n)))))

(define/contract (basin-size b)
  (-> (listof point?) number?)
  (~> b
      list->set
      set->list
      length))

(~>
 (all-cords)
 (map low-point? _)
 (filter identity _)
 (map expand-basin _)
 (sort < #:key basin-size)
 reverse
 (take _ 3)
 (map basin-size _)
 (apply * _))