blob: 5b4fbb6850fd273e17551c333d8fa1fccd3f00fb (
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
|
module Main where
import Control.Monad
import qualified Data.Array as A
import Data.Function (on)
import Data.List
import qualified Data.Map.Strict as Map
import Data.Maybe
import qualified Data.Set as Set
import Input
swap :: (a, b) -> (b, a)
swap (a, b) = (b, a)
uniq :: Eq a => [a] -> [a]
uniq (x:y:zs) | x == y = uniq (x:zs)
| otherwise = x : uniq (y:zs)
uniq l = l
parse :: String -> (String, String)
parse s = let (pre, ')' : post) = break (== ')') s
in (pre, post)
main :: IO ()
main = do
input <- getInput 6
let downedges = map parse input
upedges = map swap downedges
nodes = uniq $ sort $ map fst $ downedges ++ upedges
when (any ((/= 1) . length) (groupBy ((==) `on` fst) upedges)) $
error "Input graph is not a tree"
let parentMap = Map.fromList upedges
node2index = Map.fromList (zip nodes [0..])
depthArray = A.listArray (0, length nodes - 1)
[case Map.lookup node parentMap of
Nothing -> 0
Just par -> depthArray A.! (node2index Map.! par) + 1
| node <- nodes]
checksum = sum [d | d <- A.elems depthArray] :: Int
print checksum
let src = parentMap Map.! "YOU"
dest = parentMap Map.! "SAN"
childMap = Map.fromList (map ((,) <$> fst . head <*> map snd)
$ groupBy ((==) `on` fst)
$ sortOn fst downedges)
envOf n = maybe [] pure (Map.lookup n parentMap) ++ fromMaybe [] (Map.lookup n childMap)
expand d seen border mp
| Set.null border = mp
| otherwise = let newnodes = Set.fromList (concatMap envOf border) Set.\\ seen
in expand (d + 1)
(seen <> newnodes)
newnodes
(foldr (\n -> Map.insert n (d + 1)) mp (Set.toList newnodes))
distMap = expand (0 :: Int) (Set.singleton src) (Set.singleton src) (Map.singleton src 0)
print (distMap Map.! dest)
|
