module Data.BellmanFord where
import qualified Data.Map as M
import Data.Maybe
import Control.Monad
import Data.List.Extras.Argmax
relaxEdge nodeWeights edgeWeights x (k, weight) = argmin snd ((k, weight) : map (\y -> (y, maybe (1/0) (\z -> snd (fromJust (M.lookup y nodeWeights)) + z) (M.lookup (y, x) edgeWeights))) (M.keys nodeWeights))
bellmanFord :: (Ord a) => M.Map (a, a) Double -> a -> M.Map a (a, Double)
bellmanFord gr x = foldl (\weights _ -> M.mapWithKey (relaxEdge weights gr) weights) weights [1..M.size gr+1] where
keys = map fst (M.keys gr) ++ map snd (M.keys gr)
weights = M.insert x (x, 0) $ M.fromList $ map (\x -> (x, (x, 1/0))) keys
retrievePath :: (Ord a) => M.Map a (a, Double) -> a -> a -> [a]
retrievePath mp a a2 | a == a2 = [a]
retrievePath mp a a2 = retrievePath mp a (fst (fromJust (M.lookup a2 mp))) ++ [a2]
cycles :: (Ord a) => M.Map (a, a) () -> a -> Maybe a
cycles gr x = do
(y, z) <- M.lookup x weights
guard (round z < 0)
return y where
gr' = fmap (const (1)) gr
weights = bellmanFord gr' x