43c3c7311a
it wooooorks this kinda stuff is extremely hard in Haskell, and that Haskell programmers on the internet all seem to be galaxy brained wizards isn’t exactly helpful either (looking at you, Data.List.permutations)
87 lines
2.6 KiB
Haskell
87 lines
2.6 KiB
Haskell
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
|