AoC/2022/day16/day16.hs

import Control.Lens
import Data.Array
import Data.Int
import Data.List
import Data.List.Split
import Data.Maybe

type Valve = (Int8, [Int8])

type DMatrix = Array (Int8, Int8) Int8

main :: IO ()
main = do
  input <- map words . lines <$> readFile "input"
  let start = fromIntegral $ fromJust $ elemIndex "AA" $ map (!! 1) input
  let v = valveNum input
  let distanceMatrix =
        dijkstraAll (buildMatrix v) v (pred $ fromIntegral $ length v)
  print $ part1 distanceMatrix v start
  print v

valveNum :: [[String]] -> [Valve]
valveNum input = over (traverse . _2 . traverse) switch $ map snd parsed
  where
    parsed = map parse input
    valves = map fst parsed
    switch x = fromIntegral $ fromJust $ elemIndex x valves
    parse (_:name:_:_:rate:_:_:_:_:ts) = (name, (r, nexts))
      where
        r = read $ init $ last $ splitOn "=" rate
        nexts = map (\(x:y:_) -> [x, y]) ts

buildMatrix :: [Valve] -> DMatrix
buildMatrix valves = n
  where
    m =
      array
        ((0, 0), (l, l))
        [((x, y), maxBound :: Int8) | x <- [0 .. l], y <- [0 .. l]]
    n = m // [((i, i), 0) | i <- [0 .. l]]
    l = pred $ fromIntegral $ length valves

dijkstra :: DMatrix -> [Valve] -> [(Int8, Int8)] -> Int8 -> DMatrix
dijkstra m vs ((x, d):xs) n =
  dijkstra newM newVs (xs ++ zip nexts (repeat $ succ d)) n
  where
    newM =
      m //
      [ let d' = minimum [m ! (n, i), d]
         in ((n, i), d')
      | i <- nexts
      ]
    newVs = vs & element x' .~ (0, [])
    nexts = snd (vs !! x')
    x' = fromIntegral x
dijkstra m _ [] _ = m

dijkstraAll :: DMatrix -> [Valve] -> Int8 -> DMatrix
dijkstraAll m vs 0 = dijkstra m vs [(0, 1)] 0
dijkstraAll m vs x = dijkstraAll (dijkstra m vs [(x, 1)] x) vs (pred x)

perm30 :: DMatrix -> [(Int8, Int8)] -> Int -> Int8 -> [[(Int8, Int8)]]
perm30 _ [] _ _ = [[]]
perm30 m xxs n current
  | n > 30 = [[]]
  | otherwise =
    [y : ys | (y, xs) <- picks xxs, ys <- perm30 m xs (succ n + dist y) (fst y)]
  where
    picks (x:xs) = (x, xs) : [(y, x : ys) | (y, ys) <- picks xs]
    picks [] = []
    dist y = fromIntegral (m ! (current, fst y))

part1 :: DMatrix -> [Valve] -> Int8 -> Int
part1 m valves start = maximum $ map (calcRoute m start 30 0) routes
  where
    valves' = filter (\(_, x) -> x /= 0) $ zip [0 ..] $ map fst valves
    routes = perm30 m valves' 0 start

calcRoute :: DMatrix -> Int8 -> Int8 -> Int -> [(Int8, Int8)] -> Int
calcRoute m current time acc (node:nodes)
  | newTime > 0 = calcRoute m (fst node) newTime newAcc nodes
  | otherwise = acc
  where
    newTime = pred $ time - m ! (current, fst node)
    newAcc = acc + fromIntegral (snd node) * fromIntegral newTime
calcRoute _ _ _ acc [] = acc