{-# LANGUAGE CPP, BangPatterns #-}
{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}
module Main ( main, test, maxDistances ) where
import System.Environment
import Data.Array.Repa
-- <<Graph
type Weight = Int
type Graph r = Array r DIM2 Weight
-- >>
-- -----------------------------------------------------------------------------
-- shortestPaths
-- <<shortestPaths
shortestPaths :: Graph U -> Graph U
shortestPaths g0 = go g0 0 -- <2>
where
Z :. _ :. n = extent g0 -- <1>
go !g !k | k == n = g -- <3>
| otherwise =
let g' = computeS (fromFunction (Z:.n:.n) sp) -- <4>
in go g' (k+1) -- <5>
where
sp (Z:.i:.j) = min (g ! (Z:.i:.j)) (g ! (Z:.i:.k) + g ! (Z:.k:.j)) -- <6>
-- >>
-- -----------------------------------------------------------------------------
-- <<maxDistance
maxDistance :: Weight -> Weight -> Weight
maxDistance x y
| x == inf = y
| y == inf = x
| otherwise = max x y
-- >>
maxDistances :: Graph U -> Array U DIM1 Weight
maxDistances = foldS maxDistance inf
-- -----------------------------------------------------------------------------
-- Testing
-- <<inf
inf :: Weight
inf = 999
-- >>
testGraph :: Graph U
testGraph = toAdjMatrix $
[[ 0, inf, inf, 13, inf, inf],
[inf, 0, inf, inf, 4, 9],
[ 11, inf, 0, inf, inf, inf],
[inf, 3, inf, 0, inf, 7],
[ 15, 5, inf, 1, 0, inf],
[ 11, inf, inf, 14, inf, 0]]
-- correct result:
expectedResult :: Graph U
expectedResult = toAdjMatrix $
[[0, 16, inf, 13, 20, 20],
[19, 0, inf, 5, 4, 9],
[11, 27, 0, 24, 31, 31],
[18, 3, inf, 0, 7, 7],
[15, 4, inf, 1, 0, 8],
[11, 17, inf, 14, 21, 0] ]
test :: Bool
test = shortestPaths testGraph == expectedResult
toAdjMatrix :: [[Weight]] -> Graph U
toAdjMatrix xs = fromListUnboxed (Z :. k :. k) (concat xs)
where k = length xs
main :: IO ()
main = do
[n] <- fmap (fmap read) getArgs
let g = fromListUnboxed (Z:.n:.n) [0..n^(2::Int)-1] :: Graph U
print (sumAllS (shortestPaths g))