Week 8 with @yuliaju

Week8.hs

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