diff --git a/7.hs b/7.hs
index cd98c6a..ba80eab 100644
--- a/7.hs
+++ b/7.hs
@@ -1,20 +1,41 @@
 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
+type Weight   = Int
+type Program  = String
 type Programs = [Program]
 
 discardEmpty :: [String] -> [String]
 discardEmpty [""] = []
-discardEmpty xs = xs
+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
+    let nameAndWeight : programsString : _ = splitOn ")"  line
+        name          : weight         : _ = splitOn " (" nameAndWeight
         programs = discardEmpty $ splitOn ", " $ last $ splitOn " -> " programsString
-        name : weight : _ = splitOn " (" nameAndWeight
     in (name, (read weight, programs))
 
 getBottom :: Map Program (Weight, Programs) -> [(Program, (Weight, Programs))] -> Program
@@ -24,49 +45,35 @@ getBottom m l = bottomName
             (\(name, (_, programs)) set -> 
                 case programs of
                     [] -> delete name set
-                    ps -> foldr delete set ps) 
+                    ps -> foldr delete set ps)
             m l
 
 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 }
-
-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)
+    let newForest        = map cumulate forest
+        cumulativeWeight = (+ root) . sum $ getWeights newForest
+    in  Node { rootLabel = (root, cumulativeWeight), subForest = newForest }
 
+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
     input <- readFile "7.txt"
     let programsList = map parseLine $ lines input
-        programsMap = fromList programsList
-        bottomName = getBottom programsMap programsList
-        programsTree = mapToTree programsMap bottomName
-        cumulateTree = cumulate programsTree
-        unbalanced = findUnbalanced cumulateTree
+        programsMap  = fromList  programsList
+        bottomName   = getBottom programsMap programsList
+        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
\ No newline at end of file
+    print $ balanced
+    -- writeFile "7_tree.txt"     $ drawTree $ fmap show programsTree
+    -- writeFile "7_cumulate.txt" $ drawTree $ fmap show cumulateTree
\ No newline at end of file