ADAA2 Week 4 (Bellman-Ford + Dijkstra)

BellmanFord.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnicodeSyntax       #-}
module Algo.BellmanFord (findShortestPaths) where

import           Control.Monad         (liftM2)
import           Data.Foldable         (foldMap)
import           Data.Functor.Identity (runIdentity, Identity(..))
import qualified Data.HashMap.Strict   as M
import           Data.Maybe            (mapMaybe)
import           Data.Monoid           (Endo (..), appEndo)

import           Algo.Graph


findShortestPaths ∷ Vertex → Graph → Maybe (M.HashMap Vertex Weight)
findShortestPaths vertex γ =
  if redund_it == last_it
    then Just . remove_unreachable $ last_it
    else Nothing

 where

  last_it = (\n -> appEndo . foldMap Endo . replicate n) (graph_size-1) iteration initial
  redund_it = iteration last_it
  initial ∷ M.HashMap Vertex (Maybe Weight)
  initial = M.fromList [ (v, w) | v ← vertices γ, let w = if v == vertex then Just 0 else Nothing ]

  graph_size ∷ Int = length $ vertices γ

  iteration ∷ M.HashMap Vertex (Maybe Weight) → M.HashMap Vertex (Maybe Weight)
  iteration μ = runIdentity $ M.traverseWithKey (\v' w' → Identity $ l v' w') μ
   where
    l ∷ Vertex → Maybe Weight → Maybe Weight
    l v' w' = foldl min' w' [ liftM2 (+) (μ M.! s) (Just c) | (s,c) ← ngbs v' ]
    min' Nothing x = x
    min' x Nothing = x
    min' x y = min x y
    ngbs v = case M.lookup v (Algo.Graph.reverse γ) of
               Nothing → []
               Just x → M.toList x


remove_unreachable ∷ M.HashMap Vertex (Maybe Weight) → M.HashMap Vertex Weight
remove_unreachable = M.fromList . mapMaybe (uncurry (fmap . (,))) . M.toList

BF.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnicodeSyntax       #-}
import           Control.Applicative       (many, (<$>), (<*))
import           Control.Lens
import           Data.Attoparsec.Text.Lazy
import qualified Data.HashMap.Strict       as M
import           Data.List                 (nub, sort)
import           Data.Text.Lazy.IO         as T

import qualified Algo.BellmanFord          as BF
import qualified Algo.Dijkstra             as D
import           Algo.Graph

main ∷ IO ()
main = do
  mins ← mapM get_min_weight [ "data/g1.txt", "data/g2.txt", "data/g3.txt" ]
  case minimumOf (folded . folded) mins of
    Just x → print x
    Nothing → print "NULL"

get_min_weight ∷ FilePath → IO (Maybe Weight)
get_min_weight φ = do
  triples ← readTriples φ
  let
      extra_vertice ∷ Int = head [ h | h ← [1..], h `notElem` vertices' ]
      vertices' ∷ [Int] = nub . sort $ concatMap ((\(a,b) → [a,b]) . fst) triples
      appended_graph = triplesToGraph $ triples ++ [ ((extra_vertice, v),0) | v ← vertices' ]
  case BF.findShortestPaths extra_vertice appended_graph of
    Nothing → return Nothing
    Just p → do
      let reweighted_triples = [ ((v1,v2),e + (p M.! v1) - (p M.! v2)) | ((v1,v2),e) ← triples ]
          g = triplesToGraph reweighted_triples
      let reduce_weights v = map (\(dest,w) → w - p M.! v + p M.! dest) . M.toList
      return . Just . minimum $ concatMap (\v → reduce_weights v $ D.findShortestPaths g v) vertices'


readTriples ∷ FilePath → IO [((Int, Int), Int)]
readTriples φ = handle . parse parser <$> T.readFile φ
  where
    handle = either (error "jobs parser failed") id . eitherResult

parser = do
  _ ← decimal <* many space
  edges ← decimal <* many space
  count edges triple
 where
  triple = do
    v1 ← decimal <* many space
    v2 ← decimal <* many space
    w ← signed decimal <* many space
    return ((v1,v2),w)

Dijkstra.hs

{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE OverlappingInstances #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UnicodeSyntax        #-}
module Algo.Dijkstra (findShortestPaths) where

import           Control.Monad        (when)
import           Control.Monad.Reader (Reader, ask, runReader)
import           Control.Monad.State  (StateT, execStateT, get, put)
import qualified Data.HashMap.Strict  as M
import qualified Data.Heap            as H
import           Data.List            (nub, sort)
import qualified Data.Set             as S

import           Algo.Graph


type WeightMap = M.HashMap Vertex Weight
type WeightHeap = H.Heap (Maybe Weight, Vertex)
type VertexSet = S.Set Vertex

instance Ord (Maybe Weight) where
  compare (Just x) (Just y) = compare x y
  compare Nothing _ = GT
  compare _ Nothing = LT

findShortestPaths ∷ Graph → Vertex → WeightMap
findShortestPaths γ src_vertex = fst $ runReader (execStateT visite_all_vertices (d,(h,x))) (γ,vs)
 where
  -- distances
  d ∷ WeightMap = M.empty
  -- heap
  h ∷ WeightHeap = H.fromList [ (w,v) | v ← vertices γ, let w = if v == src_vertex then Just 0 else Nothing ]
  -- visited vertices
  x ∷ VertexSet = S.empty
  -- all vertices
  vs ∷ VertexSet = S.fromList . nub . sort $ src_nodes ++ dest_nodes
   where
     src_nodes ∷ [Vertex] = M.keys γ
     dest_nodes ∷ [Vertex] = concatMap M.keys $ M.elems γ

visite_all_vertices ∷ StateT (WeightMap,(WeightHeap,VertexSet)) (Reader (Graph,VertexSet)) ()
visite_all_vertices = do
  r ← visite_one_more_vertex
  when r visite_all_vertices

visite_one_more_vertex ∷ StateT (WeightMap,(WeightHeap,VertexSet)) (Reader (Graph,VertexSet)) Bool
visite_one_more_vertex = do
  (γ,vs) ← ask
  (distances,(heap,visited)) ← get
  let (maybe_path_to_new_node ∷ Maybe Weight, new_node ∷ Vertex) = H.minimum heap
  case maybe_path_to_new_node of
    Just path_to_new_node → do
      let
          distances' ∷ WeightMap = M.insert new_node path_to_new_node distances
          visited' ∷ VertexSet = S.insert new_node visited
          heap' ∷ WeightHeap = H.map update_ngbs $ H.deleteMin heap

          weight v w = (γ M.! v) M.! w
          a `connected_with` b = M.member a γ && M.member b (γ M.! a)
          update_ngbs (w, v) =
            let new_value = Just $ weight new_node v + path_to_new_node
            in  if new_node `connected_with` v && new_value < w
                  then (new_value,v)
                  else (w,v)
      put (distances', (heap',visited'))
      return $ visited' /= vs
    Nothing →
      return False

Graph.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnicodeSyntax       #-}
module Algo.Graph where

import           Control.Applicative       (many, (<$>), (<*))
import           Data.Attoparsec.Text.Lazy
import qualified Data.HashMap.Strict       as M
import           Data.List                 (nub, sort)
import           Data.Text.Lazy.IO         as T
import           Text.Printf               (printf)

type Vertex = Int
type Weight = Int
type Graph = M.HashMap Vertex (M.HashMap Vertex Weight)


triplesToGraph ∷ [((Vertex,Vertex),Weight)] → Graph
triplesToGraph list = M.fromList
  [ (v, M.fromList edges)
    | v ← nub $ sort $ map (fst . fst) list
    , let edges = map (\((_,b),c) → (b,c)) $ filter ((== v) . fst . fst) list
  ]

graphToTriples ∷ Graph → [((Vertex,Vertex),Weight)]
graphToTriples g = [ ((v1,v2),w) | v1 ← M.keys g, (v2,w) ← M.toList $ g M.! v1 ]

triplesFromFile ∷ FilePath → IO [((Vertex,Vertex),Weight)]
triplesFromFile φ = handle . parse parser <$> T.readFile φ where
  handle = either (error "graph parsed failed") id . eitherResult
  parser = do
    _ ← decimal <* many space
    edges ← decimal <* many space
    count edges triple
   where
    triple = do
      v1 ← decimal <* many space
      v2 ← decimal <* many space
      w ← signed decimal <* many space
      return ((v1,v2),w)

fromFile ∷ FilePath → IO Graph
fromFile φ = triplesToGraph <$> triplesFromFile φ

toFile ∷ FilePath → Graph → IO ()
toFile φ γ = Prelude.writeFile φ $ unlines $
  printf "%d %d" (length $ vertices γ) (length $ map (\((v1,v2),_) → (v1,v2)) $ graphToTriples γ)
  : [ printf "%d %d %d" v1 v2 w | v1 ← vertices γ, (v2,w) ← M.toList $ γ M.! v1 ]

vertices ∷ Graph → [Int]
vertices g = nub . sort $ M.keys g ++ concatMap M.keys (M.elems g)

reverse ∷ Graph → Graph
reverse = triplesToGraph . map (\((a,b),w) → ((b,a),w)) . graphToTriples