codeforces - 25 - C (hs)

cf25-c.hs

{-# LANGUAGE NewQualifiedOperators #-}
import Prelude 
import qualified Data.ByteString.Char8 as B
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector as V
import Data.List
import Control.Monad
import Control.Applicative
import Data.Maybe(fromJust)


type Graph = V.Vector (U.Vector Int)

build :: Int -> [[Int]] -> Graph
build n adj = V.fromList $ map U.fromList $ adj

relax ::  Int -> Graph -> (Int, Int, Int) -> Graph
relax n g (u, v, c) = V.imap (\i -> U.imap (min . f i)) g
    where 
        au = V.(!) g u
        av = V.(!) g v
        f i j = min (U.(!) au i + c + U.(!) av j) (U.(!) av i + c + U.(!) au j) 

combine :: Int -> (Graph, V.Vector Int) -> (Int, Int, Int) -> (Graph, V.Vector Int)
combine n (g, r) (u, v, c) = (g', V.cons s r)
    where 
        g' = relax n g (u, v, c)
        s  = V.sum $ V.map U.sum g'

solve ::  Int -> [[Int]] -> [(Int, Int, Int)] -> [Int]
solve n adj edges = V.toList $ snd $ V.foldl' (combine n) (build n adj, V.empty) (V.fromList edges)

ints :: IO [Int]
ints = map (fst . fromJust . B.readInt) . B.words <$> B.getLine

main = do
    [n] <- ints
    adj <- replicateM n ints
    [k] <- ints
    edges <- replicateM k $ (\(a:b:c:_) -> (a-1, b-1, c)) <$> ints
    let r = reverse $ solve n adj edges
    putStrLn $ intercalate " " $ map show r