bf.hs

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}

import Util
import Data.Ord
import Data.List
import System.IO.Unsafe
import System.Mem.StableName
import qualified Data.IntMap.Strict as M

type Vertex = Int
type Weight = Int
type Edge = (Vertex, Vertex, Weight)

parseGragh = map parse . tail . lines
  where parse s = let [v1, v2, w] = map read $ words s
                  in (v1, v2, w)

data Distance = Inf | Only Weight deriving (Show, Eq)

instance Num Distance where
  (+) Inf a = Inf
  (+) a Inf = Inf
  (+) (Only a) (Only b) = Only (a + b)
  fromInteger a = Only (fromInteger a)

instance Ord Distance where
  compare Inf a = GT
  compare a Inf = LT
  compare (Only a) (Only b) = compare a b

class DistanceMap m where
  get :: m -> Vertex -> Distance
  assoc :: m -> Vertex -> Distance -> m

type DM = M.IntMap Weight

instance DistanceMap DM where
  get m v = case M.lookup v m of
              Just x -> Only x
              Nothing -> Inf
  assoc m v Inf = M.delete v m
  assoc m v (Only x) = M.insert v x m

newtype PotMap = PotMap DM deriving (Show, Eq)

instance DistanceMap PotMap where
  get m v = min 0 (get m v)
  assoc (PotMap m) v d = PotMap (assoc m v d)

identical !a !b = unsafePerformIO io
  where
    io = do
           aRef <- makeStableName a
           bRef <- makeStableName b
           return (aRef == bRef)

bellmanFord :: (DistanceMap m) => m -> [Edge] -> Maybe m
bellmanFord m edges = loop m 0 (length edges)
  where
    loop dists step size =
      if step > size then
        Nothing
      else
        let newDists = foldl' proc dists edges
        in
          if identical newDists dists then
            Just dists
          else
            loop newDists (step + 1) size

    proc dists (v1, v2, w) =
      let d = get dists v2
          d' = get dists v1 + Only w
      in
        if d' < d then
          assoc dists v2 d'
        else
          dists

main = runP parseGragh compute
  where
    compute g = bellmanFord (PotMap M.empty) g