Paste: AOC 2021, day 21

USING: arrays combinators.extras formatting generalizations
kernel math math.order math.ranges math.vectors prettyprint
sequences ;
IN: 2021.21

CONSTANT: start1 10
CONSTANT: start2 6

: roll ( start -- end move )
    [ dup 100 mod 1 + ] thrice
    [ + ] [ + ] [ swap ] tri* ;

: pos ( pos move -- pos' ) + 1 - 10 mod 1 + ;

:: part1 ( -- score rolls )
    start1 start2 0 0 0 :> ( pos1! pos2! score1! score2! rolls! )
    1 [
        roll rolls 3 + rolls!
        pos1 pos pos1!
        score1 pos1 + score1!
        score1 1000 >= [ f ] [
            roll rolls 3 + rolls!
            pos2 pos pos2!
            score2 pos2 + score2!
            score2 1000 >= not
        ] if
    ] loop drop score1 score2 min rolls ;

DEFER: dirac
MEMO: turn ( score pos v die-sum -- v )
    [ [ [ v* supremum ] [ pos ] bi* ] keepd n*v ] 2keepd
    [ reverse v* [ [ v+ ] keep ] [ v+ ] bi* ] keep
    pick over v* supremum 21 >= [ 2nip ] [ dirac ] if ;

MEMO: dirac ( score pos v -- v )
    reverse 3 9 [a,b] [ turn ] 3 nwith map
    { 1 3 6 7 6 3 1 } [ v*n ] 2map
    { 0 0 } [ v+ ] reduce ;
    
part1 * .
{ 0 0 } start1 start2 2array { 0 1 } dirac supremum .