Bellman Ford.hs

import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Sequence as Seq
import qualified Data.Array as Array
import qualified Data.Maybe as Maybe
import qualified Control.Monad.ST as ST
import qualified Control.Monad as CM
import qualified Data.PQueue.Prio.Min as MinHeap

main :: IO ()
main = do
  g1File <- readFile "./src/Course4/g3.txt"
  let g1Data@(g1Vc, _, _) = parseData g1File
      g1Result = bellmanFord (buildAdjacencyListByHead g1Data) g1Vc 1
  print g1Result

parseData :: String -> (Int, Int, [(Int, Int, Double)])
parseData s =
  (vertexCount, edgeCount, edges)
  where
    fileLines = lines s
    (vertexCount, edgeCount) = parseMetaData . head $ fileLines
    edges = List.map parseEdge . tail $ fileLines
    parseMetaData :: String -> (Int, Int)
    parseMetaData line =
      case List.map read . words $ line of
        [x,y] -> (x, y)
        _ -> error "Invalid data format"
    parseEdge :: String -> (Int, Int, Double)
    parseEdge line =
      case words line of
        (edgeTail:edgeHead:weight:_) -> (read edgeTail, read edgeHead, read weight)
        _ -> error "Invalid data format"

buildAdjacencyListByHead :: (Fractional a, Ord a) => (Int, Int, [(Int, Int, a)]) -> IntMap.IntMap [(Int, a)]
buildAdjacencyListByHead (vertexCount, _, edges) =
  List.foldl' (\acc (t, h, w) -> IntMap.update (\edges' -> Just ((t,w):edges')) h acc) initialMap $ edges
  where
    initialMap = IntMap.fromList (List.zip [1..vertexCount] (List.repeat []))
    
bellmanFord :: IntMap.IntMap [(Int, Double)] -> Int -> Int -> Either String (Array.Array Int Double)
bellmanFord g vertexCount s =
  if hasNegativeCycle then
    Left "Bellman-Ford algorithm halted: negative cycle found"
  else
    Right resultArray
  where
    memo = Array.array ((0,1), (vertexCount - 1, vertexCount)) [generator (x,y) | x <- [0..(vertexCount - 1)], y <- [1..vertexCount]]
    generator pair@(i,v) =
      if i == 0 then
        if v == s then
          (pair, 0)
        else
          (pair, 1/0)
      else
        (pair, findMin pair)
    findMin (i,v) = min ({-# SCC accessMemo #-} (Array.!) memo (i-1,v)) ({-# SCC foldl' #-} List.foldl' (\acc (t,w) -> {-# SCC foldStep #-} min acc ((Array.!) memo (i-1,t) + w)) (1/0) ({-# SCC accessGraph #-} (IntMap.!) g v))
    resultList = map (\((_,v), d) -> (v,d)) . filter (\((i,_), _) -> i == vertexCount - 1 ) . Array.assocs $ memo
    resultArray = Array.array (1, vertexCount) resultList
    hasNegativeCycle = any (\(v, d) -> any (\(t, w) -> d > resultArray Array.! t + w) (g IntMap.! v)) resultList

{-# LANGUAGE FlexibleContexts #-}
module Main where

-- import Debug.Trace

import Control.Monad
import Data.Array (Array)
import Data.Coerce
import Data.Maybe
import Data.Semigroup
import System.IO

import qualified Data.Array.IArray as IArray
import qualified Data.Array.ST.Safe as STArray
import qualified Data.ByteString.Char8 as BS
import qualified Data.IntMap.Strict as IntMap
import qualified Data.List as List

main :: IO ()
main = do
  (g1Vc, _, g1Edges) <- parseData "./data.txt"
  let g1 = buildAdjacencyListByHead g1Vc g1Edges
  let g1Result = bellmanFord g1 g1Vc 1
  print g1Result

parseData :: FilePath -> IO (Int, Int, [(Int, Int, Double)])
parseData fp = withFile fp ReadMode $ \h -> do
  (vertexCount, edgeCount) <- getMetaData h
  edges <- replicateM edgeCount (getEdge h)
  pure (vertexCount, edgeCount, edges)
  where
    getMetaData h = parseMetaData <$> BS.hGetLine h
    parseMetaData line = fromJust $ do
      (x, line') <- BS.readInt line
      (y, line'') <- BS.readInt (dropSpace line')
      if BS.null (dropSpace line'')
        then pure (x, y)
        else error "Invalid data format"

    getEdge h = parseEdge <$> BS.hGetLine h
    parseEdge line = fromJust $ do
      (edgeTail, line') <- BS.readInt line
      (edgeHead, line'') <- BS.readInt (dropSpace line')
      (weight, line''') <- BS.readInteger (dropSpace line'')
      if BS.null (dropSpace line''')
        then pure (edgeTail, edgeHead, fromInteger weight)
        else error "Invalid data format"

    dropSpace = BS.dropWhile (== ' ')

buildAdjacencyListByHead :: (Fractional a, Ord a) => Int -> [(Int, Int, a)] -> Array Int [(Int, a)]
buildAdjacencyListByHead vertexCount =
  IArray.listArray (1, vertexCount)
  . IntMap.elems
  . List.foldl' (\acc (f, t, w) -> IntMap.adjust ((f, w) :) t acc) initialMap
  where
    initialMap = IntMap.fromAscList (List.zip [1 .. vertexCount] (List.repeat []))
    
bellmanFord :: Array Int [(Int, Double)] -> Int -> Int -> Either String (Array Int Double)
bellmanFord g vC s =
  if hasNegativeCycle then
    Left "Bellman-Ford algorithm halted: negative cycle found"
  else
    Right resultArray
  where
    calculated = STArray.runSTUArray $ do
      array <- STArray.newArray (1, vC) (1 / 0)
      STArray.writeArray array s 0
      go array 1 1
      where
        go array stage v
          | v > vC =
            if stage < vC
            then go array (stage + 1) 1
            else pure array
          | otherwise = do
            forM_ (g IArray.! v) $ coreCmp array v
            go array stage (v + 1)

        coreCmp array v (f, w) = do
          vW <- STArray.readArray array v
          fW <- STArray.readArray array f
          STArray.writeArray array v (min (fW + w) vW)
        {-# INLINE coreCmp #-}

    resultList = IArray.elems calculated
    resultArray = IArray.listArray (1, vC) resultList

    getResultOf = (resultArray IArray.!)
    {-# INLINE getResultOf #-}

    hasNegativeCycle =
      coerce
      . foldMap (\(t, efs) -> foldMap (\(f, w) -> coerce (getResultOf f + w < getResultOf t) :: Any) efs)
      $ IArray.assocs g