46 lines
1.5 KiB
Scheme
46 lines
1.5 KiB
Scheme
(import :gerbil/gambit/ports)
|
|
(import :std/pregexp)
|
|
(import :std/srfi/1)
|
|
(import :std/iter)
|
|
|
|
(export main)
|
|
|
|
(defstruct node (children metadata))
|
|
|
|
(def (parse-node lst)
|
|
(match lst
|
|
([m n . rest]
|
|
(let*-values (((children rest) (parse-children rest m))
|
|
((metadata rest) (split-at rest n)))
|
|
(values (make-node children metadata) rest)))))
|
|
|
|
(def (parse-children lst m)
|
|
(if (= m 0)
|
|
(values [] lst)
|
|
(let*-values (((children rest) (parse-children lst (1- m)))
|
|
((child rest) (parse-node rest)))
|
|
(values (cons child children) rest))))
|
|
|
|
(def (dfs node f g)
|
|
(if (null? (node-children node))
|
|
(f (node-metadata node))
|
|
(g (cons (f (node-metadata node))
|
|
(map (lambda (n) (dfs n f g)) (node-children node))))))
|
|
|
|
(def (node-value node)
|
|
(if (null? (node-children node))
|
|
(apply + (node-metadata node))
|
|
(apply + (for/collect ((m (node-metadata node)))
|
|
(if (and (> m 0) (< (1- m) (length (node-children node))))
|
|
(node-value (list-ref (reverse (node-children node)) (1- m)))
|
|
0)))))
|
|
|
|
(def (main . args)
|
|
(def example-lst [2 3 0 3 10 11 12 1 1 0 1 99 2 1 1 2])
|
|
(define-values (example _) (parse-node example-lst))
|
|
(displayln (dfs example (lambda (m) (apply + m)) (lambda (m) (apply + m))))
|
|
(displayln (node-value example))
|
|
(def lst (map string->number (pregexp-split " " (call-with-input-file "input.txt" read-line))))
|
|
(define-values (tree _) (parse-node lst))
|
|
(displayln (dfs tree (lambda (m) (apply + m)) (lambda (m) (apply + m))))
|
|
(displayln (node-value tree)))
|