Day 7, part 2 - finished, not very pretty but gets it done

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(..), unfoldTree, drawTree)
 
 type Weight = Int
 type Program = String
@@ -30,6 +30,33 @@ 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 (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 }
+
+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)
+
 main :: IO ()
 main = do
     input <- readFile "7.txt"
@@ -37,5 +64,9 @@ main = do
         programsMap = fromList programsList
         bottomName = getBottom programsMap programsList
         programsTree = mapToTree programsMap bottomName
+        cumulateTree = cumulate programsTree
+        unbalanced = findUnbalanced cumulateTree
     print $ bottomName
-    writeFile "7_tree.txt" $ drawTree $ fmap show programsTree
+    writeFile "7_tree.txt" $ drawTree $ fmap show programsTree
+    writeFile "7_cumulate.txt" $ drawTree $ fmap show cumulateTree
+    print $ unbalanced