Created
October 6, 2016 15:52
-
-
Save vaibhavsagar/eb451bd87e1b841d5a0577ef9622244c to your computer and use it in GitHub Desktop.
Week 8 with @yuliaju
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import Data.List (groupBy, sortBy) | |
import Data.Ord (comparing) | |
data DisjointSet = DisjointSet {index :: Int, value :: Int, parent :: Int} deriving Show | |
find :: [DisjointSet] -> Int -> Int | |
find ls i = let | |
djs = ls !! (i - 1) | |
ix = index djs | |
p = parent djs | |
in if p == ix | |
then ix | |
else find ls p | |
replaceNth n newVal (x:xs) | |
| n == 0 = newVal:xs | |
| otherwise = x:replaceNth (n-1) newVal xs | |
union :: [DisjointSet] -> (Int, Int) -> [DisjointSet] | |
union ls (i, j) = let | |
parent = find ls i | |
child = find ls j | |
DisjointSet ix v pa = ls !! (child - 1) | |
newDjs = DisjointSet ix v parent | |
in replaceNth (ix - 1) newDjs ls | |
makeSingleton i v = DisjointSet i v i | |
treeDiff [x] = x | |
treeDiff part = abs $ head part - head (tail part) | |
dropNth _ [] = [] | |
dropNth 0 (x:xs) = xs | |
dropNth n (x:xs) = x : dropNth (n-1) xs | |
makeTreeDiff edges values = let | |
indices = [1..length values + 1] | |
forest = zipWith makeSingleton indices values | |
unioned = foldr (flip union) forest edges | |
partitioned = groupBy (\s1 s2 -> find unioned (parent s1) == find unioned (parent s2)) $ | |
sortBy (comparing (find unioned . parent)) unioned | |
summed = map (sum . map value) partitioned | |
in treeDiff summed | |
pairToTuple [x,y] = (x,y) | |
main = do | |
vertices <- fmap read getLine :: IO Int | |
values <- fmap (map read . words) getLine :: IO [Int] | |
edges <- (mapM (\_ -> fmap (pairToTuple . map read . words) getLine :: IO (Int, Int))) [2..vertices] | |
let mini = minimum $ map ((`makeTreeDiff` values) . (`dropNth` edges)) [0..vertices - 1] | |
print mini |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment