Day 7 - refactored for clarity

7.hs

@@ -1,6 +1,6 @@
 import Data.HashMap (Map, fromList, toList, delete, findWithDefault)
 import Data.List.Split
-import Data.Tree (Tree(..), unfoldTree, drawTree)
+import Data.Tree (Tree(..), Forest, unfoldTree, drawTree)
 
 type Weight   = Int
 type Program  = String
@@ -10,11 +10,32 @@ discardEmpty :: [String] -> [String]
 discardEmpty [""] = []
 discardEmpty xs   = xs
 
+minOrZero :: Int -> Int -> Int
+minOrZero x 0 = x
+minOrZero 0 y = y
+minOrZero x y = min x y
+
+getWeights :: Forest (Weight, Int) -> [Weight]
+getWeights ts = fmap (\(Node (_, weight) _) -> weight) ts
+
+findRootWithWeight :: Int -> Forest (Weight, Int) -> Weight
+findRootWithWeight w ts = foldr (\(Node (root, weight) _) acc -> if weight == w then root else acc) 0 ts
+
+-- input:  [Weight > 0] where all weights are the same except one
+-- output: (difference between the median weight and the minority weight, the minority weight)
+getDiffMinority :: [Int] -> (Int, Int)
+getDiffMinority ns =
+    let diffs = filter (/= 0) $ map (+ (- head ns)) ns
+    in  case length diffs of
+        0 -> (0, -1)
+        1 -> (- head diffs, head ns + head diffs)
+        _ -> (  head diffs, head ns)
+
 parseLine :: String -> (Program, (Weight, Programs))
 parseLine line =
     let nameAndWeight : programsString : _ = splitOn ")"  line
-        programs = discardEmpty $ splitOn ", " $ last $ splitOn " -> " programsString
         name          : weight         : _ = splitOn " (" nameAndWeight
+        programs = discardEmpty $ splitOn ", " $ last $ splitOn " -> " programsString
     in (name, (read weight, programs))
 
 getBottom :: Map Program (Weight, Programs) -> [(Program, (Weight, Programs))] -> Program
@@ -30,32 +51,18 @@ getBottom m l = bottomName
 mapToTree :: Map Program (Weight, Programs) -> Program -> Tree Weight
 mapToTree m = unfoldTree (\s -> findWithDefault undefined s m)
 
-cumulate :: Tree Int -> Tree (Int, Int)
+cumulate :: Tree Weight -> Tree (Weight, Int)
 cumulate (Node root forest) =
     let newForest        = map cumulate forest
-        cumulativeWeight = foldr (\(Node (_, subWeight) _) weight -> weight + subWeight) 0 newForest
-    in  Node { rootLabel = (root, cumulativeWeight + root), subForest = newForest }
+        cumulativeWeight = (+ root) . sum $ getWeights newForest
+    in  Node { rootLabel = (root, cumulativeWeight), subForest = newForest }
 
-findUnbalanced :: Tree (Int, Int) -> Int
-findUnbalanced (Node _ forest) =
-    let (diff, notMedian) = getDiffMedian $ fmap (\(Node (_, subWeight) _) -> subWeight) forest
-        unbalanced = diff + foldr (\(Node (subRoot, subWeight) _) acc -> if subWeight == notMedian then subRoot else acc) 0 forest
-        subUnbalanced = sum $ map findUnbalanced forest
-    in  minOrZero unbalanced subUnbalanced
-
-minOrZero :: Int -> Int -> Int
-minOrZero 0 0 = 0
-minOrZero x 0 = x
-minOrZero 0 y = y
-minOrZero x y = min x y
-
-getDiffMedian :: [Int] -> (Int, Int)
-getDiffMedian ns =
-    let diffs = filter (/= 0) $ map (+ (- head ns)) ns
-    in  case length diffs of
-        0 -> (0, -1)
-        1 -> (- head diffs, head ns + head diffs)
-        _ -> (  head diffs, head ns)
+findBalanced :: Tree (Weight, Int) -> Int
+findBalanced (Node _ forest) =
+    let (diff, minority) = getDiffMinority $ getWeights forest
+        balanced         = diff + findRootWithWeight minority forest
+        subBalanced      = sum $ map findBalanced forest
+    in  minOrZero balanced subBalanced
         
 main :: IO ()
 main = do
@@ -63,10 +70,10 @@ main = do
     let programsList = map parseLine $ lines input
         programsMap  = fromList  programsList
         bottomName   = getBottom programsMap programsList
-        programsTree = mapToTree programsMap bottomName
-        cumulateTree = cumulate programsTree
-        unbalanced = findUnbalanced cumulateTree
+        balanced     = findBalanced . cumulate $ mapToTree programsMap bottomName
+        -- programsTree = mapToTree programsMap bottomName
+        -- cumulateTree = cumulate  programsTree
     print $ bottomName
-    writeFile "7_tree.txt" $ drawTree $ fmap show programsTree
-    writeFile "7_cumulate.txt" $ drawTree $ fmap show cumulateTree
-    print $ unbalanced
+    print $ balanced
+    -- writeFile "7_tree.txt"     $ drawTree $ fmap show programsTree
+    -- writeFile "7_cumulate.txt" $ drawTree $ fmap show cumulateTree