{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Data.DiGraph.FloydWarshall
(
DenseAdjMatrix
, AdjacencySets
, fromAdjacencySets
, toAdjacencySets
, ShortestPathMatrix(..)
, floydWarshall
, shortestPath
, distance
, diameter
, distMatrix_
, floydWarshall_
, diameter_
, shortestPaths_
) where
import Control.DeepSeq
import Data.Foldable
import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet as HS
import Data.Massiv.Array as M
import GHC.Generics
import Numeric.Natural
type DenseAdjMatrix = Array U Ix2 Int
type AdjacencySets = HM.HashMap Int (HS.HashSet Int)
fromAdjacencySets :: AdjacencySets -> DenseAdjMatrix
fromAdjacencySets :: AdjacencySets -> DenseAdjMatrix
fromAdjacencySets AdjacencySets
g = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq (forall ix. Index ix => ix -> Sz ix
Sz (Int
n Int -> Int -> Ix2
:. Int
n)) forall {a}. Num a => Ix2 -> a
go
where
n :: Int
n = forall k v. HashMap k v -> Int
HM.size AdjacencySets
g
go :: Ix2 -> a
go (Int
i :. Int
j)
| (Int, Int) -> Bool
isEdge (Int
i, Int
j) = a
1
| Bool
otherwise = a
0
isEdge :: (Int, Int) -> Bool
isEdge (Int
a, Int
b) = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
HS.member Int
b) forall a b. (a -> b) -> a -> b
$ forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup Int
a AdjacencySets
g
toAdjacencySets :: DenseAdjMatrix -> AdjacencySets
toAdjacencySets :: DenseAdjMatrix -> AdjacencySets
toAdjacencySets = forall ix r e a.
(Index ix, Source r e) =>
(a -> ix -> e -> a) -> a -> Array r ix e -> a
ifoldlS forall {a}.
(Eq a, Num a) =>
AdjacencySets -> Ix2 -> a -> AdjacencySets
f forall a. Monoid a => a
mempty
where
f :: AdjacencySets -> Ix2 -> a -> AdjacencySets
f AdjacencySets
a (Int
i :. Int
j) a
x
| a
x forall a. Eq a => a -> a -> Bool
== a
0 = AdjacencySets
a
| Bool
otherwise = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HM.insertWith forall a. Semigroup a => a -> a -> a
(<>) Int
i (forall a. Hashable a => a -> HashSet a
HS.singleton Int
j) AdjacencySets
a
newtype ShortestPathMatrix = ShortestPathMatrix (Array U Ix2 (Double, Int))
deriving (Int -> ShortestPathMatrix -> ShowS
[ShortestPathMatrix] -> ShowS
ShortestPathMatrix -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ShortestPathMatrix] -> ShowS
$cshowList :: [ShortestPathMatrix] -> ShowS
show :: ShortestPathMatrix -> String
$cshow :: ShortestPathMatrix -> String
showsPrec :: Int -> ShortestPathMatrix -> ShowS
$cshowsPrec :: Int -> ShortestPathMatrix -> ShowS
Show, ShortestPathMatrix -> ShortestPathMatrix -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c/= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
== :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c== :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
Eq, Eq ShortestPathMatrix
ShortestPathMatrix -> ShortestPathMatrix -> Bool
ShortestPathMatrix -> ShortestPathMatrix -> Ordering
ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
$cmin :: ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
max :: ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
$cmax :: ShortestPathMatrix -> ShortestPathMatrix -> ShortestPathMatrix
>= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c>= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
> :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c> :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
<= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c<= :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
< :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
$c< :: ShortestPathMatrix -> ShortestPathMatrix -> Bool
compare :: ShortestPathMatrix -> ShortestPathMatrix -> Ordering
$ccompare :: ShortestPathMatrix -> ShortestPathMatrix -> Ordering
Ord, forall x. Rep ShortestPathMatrix x -> ShortestPathMatrix
forall x. ShortestPathMatrix -> Rep ShortestPathMatrix x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep ShortestPathMatrix x -> ShortestPathMatrix
$cfrom :: forall x. ShortestPathMatrix -> Rep ShortestPathMatrix x
Generic)
deriving newtype (ShortestPathMatrix -> ()
forall a. (a -> ()) -> NFData a
rnf :: ShortestPathMatrix -> ()
$crnf :: ShortestPathMatrix -> ()
NFData)
floydWarshall :: Unbox a => Real a => Array U Ix2 a -> ShortestPathMatrix
floydWarshall :: forall a. (Unbox a, Real a) => Array U Ix2 a -> ShortestPathMatrix
floydWarshall = Array U Ix2 (Double, Int) -> ShortestPathMatrix
ShortestPathMatrix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array U Ix2 (Double, Int) -> Array U Ix2 (Double, Int)
floydWarshallInternal
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
computeAs U
U
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a r.
(Real a, Source r a) =>
Array r Ix2 a -> Array D Ix2 (Double, Int)
intDistMatrix
shortestPath
:: ShortestPathMatrix
-> Int
-> Int
-> Maybe [Int]
shortestPath :: ShortestPathMatrix -> Int -> Int -> Maybe [Int]
shortestPath (ShortestPathMatrix Array U Ix2 (Double, Int)
m) Int
src Int
trg
| forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
M.isEmpty DenseAdjMatrix
mat = forall a. Maybe a
Nothing
| Bool -> Bool
not (forall ix. Index ix => Sz ix -> ix -> Bool
M.isSafeIndex (forall r ix e. Size r => Array r ix e -> Sz ix
size Array U Ix2 (Double, Int)
m) (Int
src Int -> Int -> Ix2
:. Int
trg)) = forall a. Maybe a
Nothing
| (DenseAdjMatrix
mat forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
src Int -> Int -> Ix2
:. Int
trg)) forall a. Eq a => a -> a -> Bool
== (-Int
1) = forall a. Maybe a
Nothing
| Bool
otherwise = Int -> Int -> Maybe [Int]
go Int
src Int
trg
where
mat :: DenseAdjMatrix
mat = forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
M.computeAs U
U forall a b. (a -> b) -> a -> b
$ forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
M.map forall a b. (a, b) -> b
snd Array U Ix2 (Double, Int)
m
go :: Int -> Int -> Maybe [Int]
go Int
a Int
b
| Int
a forall a. Eq a => a -> a -> Bool
== Int
b = forall (m :: * -> *) a. Monad m => a -> m a
return []
| Bool
otherwise = do
Int
n <- forall ix r e.
(Index ix, Manifest r e) =>
Array r ix e -> ix -> Maybe e
M.index DenseAdjMatrix
mat (Int
a Int -> Int -> Ix2
:. Int
b)
(:) Int
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Maybe [Int]
go Int
n Int
b
distance :: ShortestPathMatrix -> Int -> Int -> Maybe Double
distance :: ShortestPathMatrix -> Int -> Int -> Maybe Double
distance (ShortestPathMatrix Array U Ix2 (Double, Int)
m) Int
src Int
trg
| forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
M.isEmpty Array U Ix2 (Double, Int)
m = forall a. Maybe a
Nothing
| Bool
otherwise = forall a. RealFrac a => a -> Maybe a
toDistance forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall ix r e.
(Index ix, Manifest r e) =>
Array r ix e -> ix -> Maybe e
M.index Array U Ix2 (Double, Int)
m (Int
src Int -> Int -> Ix2
:. Int
trg)
diameter :: ShortestPathMatrix -> Maybe Double
diameter :: ShortestPathMatrix -> Maybe Double
diameter (ShortestPathMatrix Array U Ix2 (Double, Int)
m)
| forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
M.isEmpty Array U Ix2 (Double, Int)
m = forall a. a -> Maybe a
Just Double
0
| Bool
otherwise = forall a. RealFrac a => a -> Maybe a
toDistance forall a b. (a -> b) -> a -> b
$ forall r ix e.
(HasCallStack, Shape r ix, Source r e, Ord e) =>
Array r ix e -> e
maximum' forall a b. (a -> b) -> a -> b
$ forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
M.map forall a b. (a, b) -> a
fst Array U Ix2 (Double, Int)
m
toDistance :: RealFrac a => a -> Maybe a
toDistance :: forall a. RealFrac a => a -> Maybe a
toDistance a
x
| a
x forall a. Eq a => a -> a -> Bool
== a
1forall a. Fractional a => a -> a -> a
/a
0 = forall a. Maybe a
Nothing
| Bool
otherwise = forall a. a -> Maybe a
Just a
x
intDistMatrix
:: Real a
#if MIN_VERSION_massiv(1,0,0)
=> Source r a
#else
=> Source r Ix2 a
#endif
=> Array r Ix2 a
-> Array M.D Ix2 (Double, Int)
intDistMatrix :: forall a r.
(Real a, Source r a) =>
Array r Ix2 a -> Array D Ix2 (Double, Int)
intDistMatrix = forall r ix e a.
(Index ix, Source r e) =>
(ix -> e -> a) -> Array r ix e -> Array D ix a
M.imap forall {a} {a}. (Real a, Fractional a) => Ix2 -> a -> (a, Int)
go
where
go :: Ix2 -> a -> (a, Int)
go (Int
x :. Int
y) a
e
| Int
x forall a. Eq a => a -> a -> Bool
== Int
y = (a
0, Int
y)
| a
e forall a. Ord a => a -> a -> Bool
> a
0 = (forall a b. (Real a, Fractional b) => a -> b
realToFrac a
e, Int
y)
| Bool
otherwise = (a
1forall a. Fractional a => a -> a -> a
/a
0, -Int
1)
floydWarshallInternal
:: Array U Ix2 (Double, Int)
-> Array U Ix2 (Double,Int)
floydWarshallInternal :: Array U Ix2 (Double, Int) -> Array U Ix2 (Double, Int)
floydWarshallInternal Array U Ix2 (Double, Int)
a = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Array U Ix2 (Double, Int) -> Int -> Array U Ix2 (Double, Int)
go Array U Ix2 (Double, Int)
a [Int
0..Int
nforall a. Num a => a -> a -> a
-Int
1]
where
Sz (Int
n :. Int
_) = forall r ix e. Size r => Array r ix e -> Sz ix
size Array U Ix2 (Double, Int)
a
go :: Array U Ix2 (Double, Int) -> Int -> Array U Ix2 (Double,Int)
go :: Array U Ix2 (Double, Int) -> Int -> Array U Ix2 (Double, Int)
go Array U Ix2 (Double, Int)
c Int
k = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq (forall ix. Index ix => ix -> Sz ix
Sz (Int
n Int -> Int -> Ix2
:. Int
n)) forall a b. (a -> b) -> a -> b
$ \(Int
x :. Int
y) ->
let
!xy :: Double
xy = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
y)
!xk :: Double
xk = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
k)
!ky :: Double
ky = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
k Int -> Int -> Ix2
:. Int
y)
!nxy :: Int
nxy = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
y)
!nxk :: Int
nxk = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$! Array U Ix2 (Double, Int)
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
k)
in if Double
xy forall a. Ord a => a -> a -> Bool
> Double
xk forall a. Num a => a -> a -> a
+ Double
ky then (Double
xk forall a. Num a => a -> a -> a
+ Double
ky, Int
nxk) else (Double
xy, Int
nxy)
distMatrix_
#if MIN_VERSION_massiv(1,0,0)
:: Source r Int
#else
:: Source r Ix2 Int
#endif
=> Array r Ix2 Int
-> Array M.D Ix2 Double
distMatrix_ :: forall r. Source r Int => Array r Ix2 Int -> Array D Ix2 Double
distMatrix_ = forall r ix e a.
(Index ix, Source r e) =>
(ix -> e -> a) -> Array r ix e -> Array D ix a
M.imap forall {a} {a}. (Real a, Fractional a) => Ix2 -> a -> a
go
where
go :: Ix2 -> a -> a
go (Int
x :. Int
y) a
e
| Int
x forall a. Eq a => a -> a -> Bool
== Int
y = a
0
| a
e forall a. Ord a => a -> a -> Bool
> a
0 = forall a b. (Real a, Fractional b) => a -> b
realToFrac a
e
| Bool
otherwise = a
1forall a. Fractional a => a -> a -> a
/a
0
floydWarshall_
:: Array U Ix2 Double
-> Array U Ix2 Double
floydWarshall_ :: Array U Ix2 Double -> Array U Ix2 Double
floydWarshall_ Array U Ix2 Double
a = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Array U Ix2 Double -> Int -> Array U Ix2 Double
go Array U Ix2 Double
a [Int
0..Int
nforall a. Num a => a -> a -> a
-Int
1]
where
Sz (Int
n :. Int
_) = forall r ix e. Size r => Array r ix e -> Sz ix
size Array U Ix2 Double
a
go :: Array U Ix2 Double -> Int -> Array U Ix2 Double
go :: Array U Ix2 Double -> Int -> Array U Ix2 Double
go Array U Ix2 Double
c Int
k = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
Seq (forall ix. Index ix => ix -> Sz ix
Sz (Int
n Int -> Int -> Ix2
:. Int
n)) forall a b. (a -> b) -> a -> b
$ \(Int
x :. Int
y) ->
let
!xy :: Double
xy = Array U Ix2 Double
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
y)
!xk :: Double
xk = Array U Ix2 Double
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
x Int -> Int -> Ix2
:. Int
k)
!ky :: Double
ky = Array U Ix2 Double
c forall r ix e.
(HasCallStack, Manifest r e, Index ix) =>
Array r ix e -> ix -> e
M.! (Int
k Int -> Int -> Ix2
:. Int
y)
in if Double
xy forall a. Ord a => a -> a -> Bool
> Double
xk forall a. Num a => a -> a -> a
+ Double
ky then Double
xk forall a. Num a => a -> a -> a
+ Double
ky else Double
xy
shortestPaths_ :: Array U Ix2 Int -> Array U Ix2 Double
shortestPaths_ :: DenseAdjMatrix -> Array U Ix2 Double
shortestPaths_ = Array U Ix2 Double -> Array U Ix2 Double
floydWarshall_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r e r' ix.
(Manifest r e, Load r' ix e) =>
r -> Array r' ix e -> Array r ix e
computeAs U
U forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r. Source r Int => Array r Ix2 Int -> Array D Ix2 Double
distMatrix_
diameter_ :: Array U Ix2 Int -> Maybe Natural
diameter_ :: DenseAdjMatrix -> Maybe Natural
diameter_ DenseAdjMatrix
g
| forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
M.isEmpty DenseAdjMatrix
g = forall a. a -> Maybe a
Just Natural
0
| Bool
otherwise = let x :: Natural
x = forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ forall r ix e.
(HasCallStack, Shape r ix, Source r e, Ord e) =>
Array r ix e -> e
maximum' forall a b. (a -> b) -> a -> b
$ DenseAdjMatrix -> Array U Ix2 Double
shortestPaths_ DenseAdjMatrix
g
in if Natural
x forall a. Eq a => a -> a -> Bool
== forall a b. (RealFrac a, Integral b) => a -> b
round (Double
1forall a. Fractional a => a -> a -> a
/Double
0 :: Double) then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just Natural
x