Day 6

1 files changed, 64 insertions, 0 deletions

diff --git a/2019/6.hs b/2019/6.hs
new file mode 100644
index 0000000..5b4fbb6
--- /dev/null
+++ b/2019/6.hs
@@ -0,0 +1,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)