76 lines
3.0 KiB
Haskell
76 lines
3.0 KiB
Haskell
import Data.HashMap (Map, fromList, toList, delete, findWithDefault)
|
|
import Data.List.Split
|
|
import Data.Tree (Tree(..), Forest, unfoldTree, drawTree)
|
|
|
|
type Weight = Int
|
|
type Program = String
|
|
type Programs = [Program]
|
|
|
|
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 = fmap (\(Node (_, weight) _) -> weight)
|
|
|
|
findRootWithWeight :: Int -> Forest (Weight, Int) -> Weight
|
|
findRootWithWeight w = foldr (\(Node (root, weight) _) acc -> if weight == w then root else acc) 0
|
|
|
|
-- 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
|
|
name : weight : _ = splitOn " (" nameAndWeight
|
|
programs = discardEmpty . splitOn ", " . last . splitOn " -> " $ programsString
|
|
in (name, (read weight, programs))
|
|
|
|
getBottom :: Map Program (Weight, Programs) -> [(Program, (Weight, Programs))] -> Program
|
|
getBottom m l = bottomName
|
|
where (bottomName, _) : _ = toList $ foldr (\(name, (_, programs)) set ->
|
|
case programs of
|
|
[] -> delete name set
|
|
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 Weight -> Tree (Weight, Int)
|
|
cumulate (Node root forest) =
|
|
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
|
|
programsList <- map parseLine . lines <$> readFile "07.txt"
|
|
let programsMap = fromList programsList
|
|
bottomName = getBottom programsMap programsList
|
|
balanced = findBalanced . cumulate $ mapToTree programsMap bottomName
|
|
-- programsTree = mapToTree programsMap bottomName
|
|
-- cumulateTree = cumulate programsTree
|
|
print $ bottomName
|
|
print $ balanced
|
|
-- writeFile "7_tree.txt" $ drawTree $ fmap show programsTree
|
|
-- writeFile "7_cumulate.txt" $ drawTree $ fmap show cumulateTree |