summaryrefslogtreecommitdiff
path: root/2021/12.hs
blob: ba902dd9ca09d7ba14e3e43ced8832ac756c0438 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
module Main where

import qualified Data.Array as A
import Data.Bifunctor (bimap)
import Data.Char
import Data.List
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Data.Tuple (swap)

import Input
import Util


main :: IO ()
main = do
    inp <- getInput 12
    let listpairs = map (toList . splitOn (== '-')) inp
        pairs = map (\[a,b] -> (a, b)) listpairs
        nodes = uniq (sort (concat listpairs))
        nodeids = Map.fromList (zip nodes [0::Int ..])
        idpairs = map (bimap (nodeids Map.!) (nodeids Map.!)) pairs
        startid = nodeids Map.! "start"
        endid = nodeids Map.! "end"
        collect l = A.listArray (0, length l - 1) l
        graph = collect
            [collect [b | (a, b) <- idpairs ++ map swap idpairs
                     , a == node]
            | node <- [0 .. length nodes - 1]]
        issmall = collect (map (all isLower) nodes)
        npathsfrom :: Int -> Set.Set Int -> Bool -> [[Int]]
        npathsfrom node seen permissive
          | node == endid = [[endid]]
          | otherwise =
              let nexts = A.elems (graph A.! node)
                  allowed n = not (issmall A.! n) || n `Set.notMember` seen
                  doubling n = n /= startid && n /= endid && not (allowed n)
                  nexts_allowed = filter allowed nexts
                  nexts_doubling = filter doubling nexts
                  seen' | issmall A.! node = Set.insert node seen
                        | otherwise = seen
              in concatMap (\n -> map (node :) (npathsfrom n seen' permissive)) nexts_allowed
                 ++ (if permissive
                         then concatMap (\n -> map (node :) (npathsfrom n seen' False)) nexts_doubling
                         else [])
    print (length $ npathsfrom startid mempty False)
    print (length $ npathsfrom startid mempty True)