Paste: AoC 2022 Day 16

! Copyright (C) 2022 Keldan Chapman.
! See http://factorcode.org/license.txt for BSD license.
USING: kernel AOC prettyprint splitting peg.ebnf multiline math.parser
       sequences assocs accessors strings sets math arrays sorting
       deques math.order tools.continuations locals sequences.generalizations
       math.vectors ranges sequences.extras heaps ;
IN: AOC.2022.16

TUPLE: node name rate connections fastest ;
: <node> ( name rate connections -- node ) V{ } clone [ node boa 0 over ] keep set-at ;

EBNF: parse-line [=[
    num = [0-9]+ => [[ dec> ]]
    name = [A-Z]+ => [[ >string ]]
    line = "Valve "~ name " has flow rate="~ num ("; tunnel" "s"? " lead" "s"? " to valve" "s"?)~ ((","? " ")~ name)+ => [[ first3 <node> ]]
]=]

: connect-all ( nodes -- map ) [ [ [ name>> ] keep ] H{ } map>assoc ] keep [ [ [ over at ] map ] change-connections drop ] each ;
: parse ( input -- assoc ) split-lines [ parse-line ] map connect-all ;

:: fastest-path ( node target -- length )
    target node fastest>> at [
        node fastest>> values supremum :> dist!
        node fastest>> [ nip dist = ] assoc-filter keys :> opts!
        [ target opts in? not ] [
            dist 1 + dist!
            dist 100 > [ target node 2array throw ] when
            opts [ connections>> ] map concat [ node fastest>> key? not ] filter opts!
            opts [ dist swap node fastest>> set-at ] each
        ] while dist
    ] unless* ;

TUPLE: path locations current opts seen ;
: <path> ( locs cur opts seen -- path ) path boa ;
: active-loc ( path -- loc ) locations>> [ second ] supremum-by ;
: potential ( item -- potential )
    [ current>> ]
    [ active-loc second 2 - 0 max 0 [a..b) <evens> ]
    [ opts>> [ rate>> ] map ]
    tri v* sum + ;

:: new-opt ( path to -- opt/f )
    path active-loc :> node
    node first2 :> ( loc time )
    loc to fastest-path 1 + :> time-cost
    time time-cost > [
        time time-cost - :> rem-time
        to rate>> rem-time * :> value
        to path opts>> remove :> rem-opts
        path seen>> to suffix :> rem-seen
        path locations>> :> locs
        node locs index locs remove-nth to rem-time 2array suffix :> rem-locs
        rem-locs path current>> value + rem-opts rem-seen <path>
    ] [ f ] if ;

: branch-item ( max item -- new-max opts )
    2dup potential > [ drop { } ] [
        [ opts>> ] [ '[ _ swap new-opt dup [ current>> '[ _ max ] dip ] when* ] map sift ]
        ! Also handle case where someone stops early
        [ clone [ [ clone ] map ] change-locations ] tri dup active-loc dup second zero? [ 2drop ] [ 0 1 rot set-nth suffix ] if
    ] if ;

: push-path ( path heap -- ) [ dup current>> ] dip heap-push ;
: pop-path ( heap -- path ) heap-pop drop ;
: peek-path ( heap -- path ) heap-peek drop ;

: branch ( max deque -- new-max heap ) tuck pop-path branch-item swapd [ over dup number? [ break ] when push-path ] each ;

: done? ( heap -- ? ) dup heap-empty? [ drop t ] [ [ heap-size 1 = ] [ peek-path ] bi opts>> length zero? and ] if ;

: initial-path ( assoc n time -- 1heap )
    '[ "AA" swap at _ swap [ _ 2array ] curry V{ } replicate-as 0 ] [ values [ rate>> 0 > ] filter [ [ rate>> ] bi@ >=< ] sort ] bi
    { } clone <path> <max-heap> [ push-path ] keep ;

: run-sim ( assoc workers time -- result ) initial-path 0 swap branch [ dup done? not ] [ branch ] while drop ;
: part-1 ( input -- result ) parse 1 30 run-sim ;
: part-2 ( input -- result ) parse 2 26 run-sim ;

MAIN: [ 16 read-day-input [ part-1 . ] [ part-2 . ] bi ]