Skip to content

Instantly share code, notes, and snippets.

@oisdk oisdk/dijkstra.hs
Last active May 7, 2019

Embed
What would you like to do?
{-# LANGUAGE DeriveFunctor, LambdaCase, FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, UndecidableInstances, RankNTypes, GeneralizedNewtypeDeriving #-}
import Data.Semigroup hiding (Sum(..))
import Control.Applicative
import Control.Monad
import Data.Functor.Classes
import Control.Monad.State
import qualified Data.Set as Set
import Data.Set (Set)
import Control.Monad.Except
import Data.List
newtype Sum a = Sum { getSum :: a } deriving (Eq, Ord, Num)
instance Show a => Show (Sum a) where
showsPrec n = showsPrec n . getSum
instance Num a => Semigroup (Sum a) where
Sum x <> Sum y = Sum (x + y)
instance Num a => Monoid (Sum a) where
mappend = (<>)
mempty = Sum 0
class (Monoid a, Semigroup a) => Group a where
inv :: a -> a
data Heap a b
= Leaf
| Node a b [Heap a b]
deriving Show
instance (Ord a, Group a) => Semigroup (Heap a b) where
Leaf <> ys = ys
xs <> Leaf = xs
Node x xv xs <> Node y yv ys
| x <= y = Node x xv (Node (inv x <> y) yv ys : xs)
| otherwise = Node y yv (Node (inv y <> x) xv xs : ys)
instance (Ord a, Group a) => Monoid (Heap a b) where
mempty = Leaf
mappend = (<>)
mconcat = mergeHeaps id
instance Functor (Heap a) where
fmap f Leaf = Leaf
fmap f (Node k x xs) = Node k (f x) (map (fmap f) xs)
(<><) :: Semigroup a => a -> Heap a b -> Heap a b
x <>< Leaf = Leaf
x <>< Node y yv ys = Node (x <> y) yv ys
mergeHeaps :: (Group a, Ord a) => (c -> Heap a b) -> [c] -> Heap a b
mergeHeaps _ [] = Leaf
mergeHeaps f (x : xs) = go x xs
where
go x [] = f x
go x1 (x2 : []) = f x1 <> f x2
go x1 (x2 : x3 : xs) = (f x1 <> f x2) <> go x3 xs
instance (Group a, Ord a) => Applicative (Heap a) where
pure x = Node mempty x []
(<*>) = ap
instance (Group a, Ord a) => Monad (Heap a) where
Leaf >>= _ = Leaf
Node k x xs >>= f = k <>< (f x <> mergeHeaps (>>= f) xs)
instance (Group a, Ord a) => Alternative (Heap a) where
empty = mempty
(<|>) = (<>)
instance (Group a, Ord a) => MonadPlus (Heap a) where
class MonadPlus m => MonadStar m where
star :: (a -> m a) -> a -> m a
plus :: (a -> m a) -> a -> m a
star f x = pure x <|> plus f x
plus f x = star f x >>= f
instance (Group a, Ord a)
=> MonadStar (Heap a) where
star f x = Node mempty x (go (f x))
where
go Leaf = []
go (Node k x xs) =
[Node k x (go (f x <> mconcat xs))]
instance MonadStar m
=> MonadStar (StateT s m) where
star f x =
StateT
(star (uncurry (runStateT . f)) . (,) x)
popMin :: (Ord a, Group a) => Heap a b -> Maybe (b, Heap a b)
popMin Leaf = Nothing
popMin (Node x xv xs) = Just (xv, x <>< mergeHeaps id xs)
toList :: (Ord a, Group a) => Heap a b -> [b]
toList = unfoldr popMin
prio :: a -> b -> Heap a b
prio k v = Node k v []
fromList :: (Ord a, Group a) => [(a,b)] -> Heap a b
fromList = foldMap (uncurry prio)
type Graph a b = b -> [(a,b)]
type GraphT a t b = b -> t (Heap a) b
instance Num a => Group (Sum a) where
inv (Sum x) = Sum (negate x)
uniq :: ( Group a
, Ord a
, Ord b
, MonadState (Set b) (t (Heap a))
, Alternative (t (Heap a)))
=> GraphT a t b -> GraphT a t b
uniq g i = do
m <- gets (Set.member i)
guard (not m)
modify (Set.insert i)
g i
path :: Functor (t (Heap a))
=> GraphT a t b
-> GraphT a t [b]
path g xs@(x:_) = fmap (:xs) (g x)
dijkstra :: (Ord a, Group a, Ord b)
=> b
-> b
-> Graph a b
-> Maybe [b]
dijkstra f t
= fmap reverse
. find ((t==) . head)
. toList
. flip evalStateT Set.empty
. flip star [f]
. path
. uniq
. fmap (lift . fromList)
graph :: Graph (Sum Int) Int
graph 1 = [( 9, 3 ),(7, 2 ),(14, 6 )]
graph 2 = [( 10, 3 ),(15, 4 ),(7, 1 )]
graph 3 = [( 2, 6 ),(11, 4 ),(9, 1 ),(10,2)]
graph 4 = [( 6, 5 ),(11, 3 ),(15, 2 )]
graph 5 = [( 9, 6 ),(6, 4 )]
graph 6 = [( 9, 5 ),(2, 3 ),(14, 1 )]
graph _ = []
-- dijkstra 1 5 graph == Just [1,3,6,5]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.