{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Graph.Base
-- Copyright   :  (c) Masahiro Sakai 2020
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-----------------------------------------------------------------------------
module ToySolver.Graph.Base
  (
  -- * Graph data types
    EdgeLabeledGraph
  , Graph
  , Vertex
  , VertexSet
  , Edge

  -- * Conversion

  -- ** Directed graphs
  , graphFromEdges
  , graphFromEdgesWith
  , graphToEdges

  -- ** Undirected graphs
  , graphFromUnorderedEdges
  , graphFromUnorderedEdgesWith
  , graphToUnorderedEdges

  -- * Operations
  , converseGraph
  , complementGraph
  , complementSimpleGraph

  -- * Properties
  , numVertexes
  , isSimpleGraph
  , isIndependentSet
  , isIndependentSetOf
  , isCliqueOf
  ) where

import Control.Monad
import Data.Array.IArray
import Data.Array.ST
import qualified Data.IntMap.Lazy as IntMap
import Data.IntMap.Lazy (IntMap)
import qualified Data.IntSet as IntSet
import Data.IntSet (IntSet)
import Data.Maybe (maybeToList)
import GHC.Stack (HasCallStack)

-- | Labelled directed graph without multiple edges
--
-- We also represent undirected graphs as symmetric directed graphs.
type EdgeLabeledGraph a = Array Vertex (IntMap a)

-- | Directed graph without multiple edges
--
-- We also represent undirected graphs as symmetric directed graphs.
type Graph = EdgeLabeledGraph ()

-- | Vertex data type
type Vertex = Int

-- | Set of vertexes
type VertexSet = IntSet

-- | Edge data type
type Edge a = (Vertex, Vertex, a)

-- | Set of edges of directed graph
graphToEdges :: EdgeLabeledGraph a -> [Edge a]
graphToEdges :: forall a. EdgeLabeledGraph a -> [Edge a]
graphToEdges EdgeLabeledGraph a
g = do
  (Vertex
node1, IntMap a
nodes) <- EdgeLabeledGraph a -> [(Vertex, IntMap a)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs EdgeLabeledGraph a
g
  (Vertex
node2, a
a) <- IntMap a -> [(Vertex, a)]
forall a. IntMap a -> [(Vertex, a)]
IntMap.toList IntMap a
nodes
  Edge a -> [Edge a]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return (Vertex
node1, Vertex
node2, a
a)

-- | Construct a directed graph from edges.
--
-- If there are multiple edges with the same starting and ending
-- vertexes, the last label is used.
graphFromEdges :: HasCallStack => Int -> [Edge a] -> EdgeLabeledGraph a
graphFromEdges :: forall a. HasCallStack => Vertex -> [Edge a] -> EdgeLabeledGraph a
graphFromEdges = (a -> a -> a) -> Vertex -> [Edge a] -> EdgeLabeledGraph a
forall a.
HasCallStack =>
(a -> a -> a) -> Vertex -> [Edge a] -> EdgeLabeledGraph a
graphFromEdgesWith a -> a -> a
forall a b. a -> b -> a
const

-- | Construct a directed graph from edges.
--
-- If there are multiple edges with the same starting and ending
-- vertexes, the labels are combined using the given function.
graphFromEdgesWith :: HasCallStack => (a -> a -> a) -> Int -> [Edge a] -> EdgeLabeledGraph a
graphFromEdgesWith :: forall a.
HasCallStack =>
(a -> a -> a) -> Vertex -> [Edge a] -> EdgeLabeledGraph a
graphFromEdgesWith a -> a -> a
_ Vertex
n [Edge a]
_ | Vertex
n Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
< Vertex
0 = [Char] -> EdgeLabeledGraph a
forall a. HasCallStack => [Char] -> a
error [Char]
"graphFromEdgesWith: number of vertexes should be non-negative"
graphFromEdgesWith a -> a -> a
f Vertex
n [Edge a]
es = (forall s. ST s (STArray s Vertex (IntMap a)))
-> EdgeLabeledGraph a
forall i e. (forall s. ST s (STArray s i e)) -> Array i e
runSTArray ((forall s. ST s (STArray s Vertex (IntMap a)))
 -> EdgeLabeledGraph a)
-> (forall s. ST s (STArray s Vertex (IntMap a)))
-> EdgeLabeledGraph a
forall a b. (a -> b) -> a -> b
$ do
  STArray s Vertex (IntMap a)
g <- (Vertex, Vertex) -> IntMap a -> ST s (STArray s Vertex (IntMap a))
forall i.
Ix i =>
(i, i) -> IntMap a -> ST s (STArray s i (IntMap a))
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Vertex
0, Vertex
nVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
-Vertex
1) IntMap a
forall a. IntMap a
IntMap.empty
  [Edge a] -> (Edge a -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Edge a]
es ((Edge a -> ST s ()) -> ST s ()) -> (Edge a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Vertex
node1, Vertex
node2, a
a) -> do
    IntMap a
m <- STArray s Vertex (IntMap a) -> Vertex -> ST s (IntMap a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Vertex (IntMap a)
g Vertex
node1
    STArray s Vertex (IntMap a) -> Vertex -> IntMap a -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Vertex (IntMap a)
g Vertex
node1 (IntMap a -> ST s ()) -> IntMap a -> ST s ()
forall a b. (a -> b) -> a -> b
$! (a -> a -> a) -> Vertex -> a -> IntMap a -> IntMap a
forall a. (a -> a -> a) -> Vertex -> a -> IntMap a -> IntMap a
IntMap.insertWith a -> a -> a
f Vertex
node2 a
a IntMap a
m
  STArray s Vertex (IntMap a) -> ST s (STArray s Vertex (IntMap a))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return STArray s Vertex (IntMap a)
g

-- | Set of edges of undirected graph represented as a symmetric directed graph.
graphToUnorderedEdges :: EdgeLabeledGraph a -> [Edge a]
graphToUnorderedEdges :: forall a. EdgeLabeledGraph a -> [Edge a]
graphToUnorderedEdges EdgeLabeledGraph a
g = do
  (Vertex
node1, IntMap a
nodes) <- EdgeLabeledGraph a -> [(Vertex, IntMap a)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs EdgeLabeledGraph a
g
  case Vertex -> IntMap a -> (IntMap a, Maybe a, IntMap a)
forall a. Vertex -> IntMap a -> (IntMap a, Maybe a, IntMap a)
IntMap.splitLookup Vertex
node1 IntMap a
nodes of
    (IntMap a
_, Maybe a
m, IntMap a
nodes2) ->
      [(Vertex
node1, Vertex
node1, a
a) | a
a <- Maybe a -> [a]
forall a. Maybe a -> [a]
maybeToList Maybe a
m] [Edge a] -> [Edge a] -> [Edge a]
forall a. [a] -> [a] -> [a]
++
      [(Vertex
node1, Vertex
node2, a
a) | (Vertex
node2, a
a) <- IntMap a -> [(Vertex, a)]
forall a. IntMap a -> [(Vertex, a)]
IntMap.toList IntMap a
nodes2]

-- | Construct a symmetric directed graph from unordered edges.
--
-- If there are multiple edges with the same starting and ending
-- vertexes, the last label is used.
graphFromUnorderedEdges :: HasCallStack => Int -> [Edge a] -> EdgeLabeledGraph a
graphFromUnorderedEdges :: forall a. HasCallStack => Vertex -> [Edge a] -> EdgeLabeledGraph a
graphFromUnorderedEdges = (a -> a -> a) -> Vertex -> [Edge a] -> EdgeLabeledGraph a
forall a.
HasCallStack =>
(a -> a -> a) -> Vertex -> [Edge a] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith a -> a -> a
forall a b. a -> b -> a
const

-- | Construct a symmetric directed graph from unordered edges.
--
-- If there are multiple edges with the same starting and ending
-- vertexes, the labels are combined using the given function.
graphFromUnorderedEdgesWith :: HasCallStack => (a -> a -> a) -> Int -> [Edge a] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith :: forall a.
HasCallStack =>
(a -> a -> a) -> Vertex -> [Edge a] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith a -> a -> a
_ Vertex
n [Edge a]
_ | Vertex
n Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
< Vertex
0 = [Char] -> EdgeLabeledGraph a
forall a. HasCallStack => [Char] -> a
error [Char]
"graphFromUnorderedEdgesWith: number of vertexes should be non-negative"
graphFromUnorderedEdgesWith a -> a -> a
f Vertex
n [Edge a]
es = (forall s. ST s (STArray s Vertex (IntMap a)))
-> EdgeLabeledGraph a
forall i e. (forall s. ST s (STArray s i e)) -> Array i e
runSTArray ((forall s. ST s (STArray s Vertex (IntMap a)))
 -> EdgeLabeledGraph a)
-> (forall s. ST s (STArray s Vertex (IntMap a)))
-> EdgeLabeledGraph a
forall a b. (a -> b) -> a -> b
$ do
  STArray s Vertex (IntMap a)
a <- (Vertex, Vertex) -> IntMap a -> ST s (STArray s Vertex (IntMap a))
forall i.
Ix i =>
(i, i) -> IntMap a -> ST s (STArray s i (IntMap a))
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Vertex
0, Vertex
nVertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
-Vertex
1) IntMap a
forall a. IntMap a
IntMap.empty
  let ins :: Vertex -> Vertex -> a -> m ()
ins Vertex
i Vertex
x a
l = do
        IntMap a
m <- STArray s Vertex (IntMap a) -> Vertex -> m (IntMap a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Vertex (IntMap a)
a Vertex
i
        STArray s Vertex (IntMap a) -> Vertex -> IntMap a -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Vertex (IntMap a)
a Vertex
i (IntMap a -> m ()) -> IntMap a -> m ()
forall a b. (a -> b) -> a -> b
$! (a -> a -> a) -> Vertex -> a -> IntMap a -> IntMap a
forall a. (a -> a -> a) -> Vertex -> a -> IntMap a -> IntMap a
IntMap.insertWith a -> a -> a
f Vertex
x a
l IntMap a
m
  [Edge a] -> (Edge a -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Edge a]
es ((Edge a -> ST s ()) -> ST s ()) -> (Edge a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Vertex
node1, Vertex
node2, a
a) -> do
    Vertex -> Vertex -> a -> ST s ()
forall {m :: * -> *}.
MArray (STArray s) (IntMap a) m =>
Vertex -> Vertex -> a -> m ()
ins Vertex
node1 Vertex
node2 a
a
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Vertex
node1 Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
node2) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ Vertex -> Vertex -> a -> ST s ()
forall {m :: * -> *}.
MArray (STArray s) (IntMap a) m =>
Vertex -> Vertex -> a -> m ()
ins Vertex
node2 Vertex
node1 a
a
  STArray s Vertex (IntMap a) -> ST s (STArray s Vertex (IntMap a))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return STArray s Vertex (IntMap a)
a

-- | Converse of a graph.
--
-- It returns another directed graph on the same set of vertices with all of the edges reversed.
-- This is also called /transpose/ or /reverse/ of a graph.
converseGraph :: EdgeLabeledGraph a -> EdgeLabeledGraph a
converseGraph :: forall a. EdgeLabeledGraph a -> EdgeLabeledGraph a
converseGraph EdgeLabeledGraph a
g = Vertex -> [Edge a] -> EdgeLabeledGraph a
forall a. HasCallStack => Vertex -> [Edge a] -> EdgeLabeledGraph a
graphFromEdges (EdgeLabeledGraph a -> Vertex
forall a. EdgeLabeledGraph a -> Vertex
numVertexes EdgeLabeledGraph a
g) [(Vertex
n2, Vertex
n1, a
l) | (Vertex
n1, Vertex
n2, a
l) <- EdgeLabeledGraph a -> [Edge a]
forall a. EdgeLabeledGraph a -> [Edge a]
graphToEdges EdgeLabeledGraph a
g]

-- | Complement of a graph
--
-- Note that applying it to a graph with no self-loops results in a graph with self-loops on all vertices.
complementGraph :: EdgeLabeledGraph a -> EdgeLabeledGraph ()
complementGraph :: forall a. EdgeLabeledGraph a -> EdgeLabeledGraph ()
complementGraph EdgeLabeledGraph a
g = (Vertex, Vertex) -> [(Vertex, IntMap ())] -> EdgeLabeledGraph ()
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (EdgeLabeledGraph a -> (Vertex, Vertex)
forall i. Ix i => Array i (IntMap a) -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds EdgeLabeledGraph a
g) [(Vertex
node, IntMap ()
toAllNodes IntMap () -> IntMap a -> IntMap ()
forall a b. IntMap a -> IntMap b -> IntMap a
IntMap.\\ IntMap a
outEdges) | (Vertex
node, IntMap a
outEdges) <- EdgeLabeledGraph a -> [(Vertex, IntMap a)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs EdgeLabeledGraph a
g]
  where
    toAllNodes :: IntMap ()
toAllNodes = [(Vertex, ())] -> IntMap ()
forall a. [(Vertex, a)] -> IntMap a
IntMap.fromAscList [(Vertex
node, ()) | Vertex
node <- EdgeLabeledGraph a -> [Vertex]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i]
indices EdgeLabeledGraph a
g]

-- | Complement of a simple graph
--
-- It ignores self-loops in the input graph and also does not add self-loops to the output graph.
complementSimpleGraph :: EdgeLabeledGraph a -> EdgeLabeledGraph ()
complementSimpleGraph :: forall a. EdgeLabeledGraph a -> EdgeLabeledGraph ()
complementSimpleGraph EdgeLabeledGraph a
g = (Vertex, Vertex) -> [(Vertex, IntMap ())] -> EdgeLabeledGraph ()
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (EdgeLabeledGraph a -> (Vertex, Vertex)
forall i. Ix i => Array i (IntMap a) -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds EdgeLabeledGraph a
g) [(Vertex
node, Vertex -> IntMap () -> IntMap ()
forall a. Vertex -> IntMap a -> IntMap a
IntMap.delete Vertex
node IntMap ()
toAllNodes IntMap () -> IntMap a -> IntMap ()
forall a b. IntMap a -> IntMap b -> IntMap a
IntMap.\\ IntMap a
outEdges) | (Vertex
node, IntMap a
outEdges) <- EdgeLabeledGraph a -> [(Vertex, IntMap a)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs EdgeLabeledGraph a
g]
  where
    toAllNodes :: IntMap ()
toAllNodes = [(Vertex, ())] -> IntMap ()
forall a. [(Vertex, a)] -> IntMap a
IntMap.fromAscList [(Vertex
node, ()) | Vertex
node <- EdgeLabeledGraph a -> [Vertex]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i]
indices EdgeLabeledGraph a
g]

-- | Number of vertexes of a graph
numVertexes :: EdgeLabeledGraph a -> Int
numVertexes :: forall a. EdgeLabeledGraph a -> Vertex
numVertexes EdgeLabeledGraph a
g =
  case EdgeLabeledGraph a -> (Vertex, Vertex)
forall i. Ix i => Array i (IntMap a) -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds EdgeLabeledGraph a
g of
    (Vertex
lb, Vertex
ub)
      | Vertex
lb Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
/= Vertex
0 -> [Char] -> Vertex
forall a. HasCallStack => [Char] -> a
error [Char]
"numVertexes: lower bound should be 0"
      | Bool
otherwise -> Vertex
ub Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1

-- | A graph is /simple/ if it contains no self-loops.
isSimpleGraph :: EdgeLabeledGraph a -> Bool
isSimpleGraph :: forall a. EdgeLabeledGraph a -> Bool
isSimpleGraph EdgeLabeledGraph a
g = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [Vertex
v Vertex -> IntMap a -> Bool
forall a. Vertex -> IntMap a -> Bool
`IntMap.notMember` IntMap a
es | (Vertex
v, IntMap a
es) <- EdgeLabeledGraph a -> [(Vertex, IntMap a)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs EdgeLabeledGraph a
g]

-- | Alias of 'isIndependentSetOf'
{-# DEPRECATED isIndependentSet "Use isIndependentSetOf instead" #-}
isIndependentSet :: EdgeLabeledGraph a -> VertexSet -> Bool
isIndependentSet :: forall a. EdgeLabeledGraph a -> VertexSet -> Bool
isIndependentSet = (VertexSet -> EdgeLabeledGraph a -> Bool)
-> EdgeLabeledGraph a -> VertexSet -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip VertexSet -> EdgeLabeledGraph a -> Bool
forall a. VertexSet -> EdgeLabeledGraph a -> Bool
isIndependentSetOf

-- | An independent set of a graph is a set of vertices such that no two vertices in the set are adjacent.
--
-- This function ignores self-loops in the input graph.
isIndependentSetOf :: VertexSet -> EdgeLabeledGraph a -> Bool
isIndependentSetOf :: forall a. VertexSet -> EdgeLabeledGraph a -> Bool
isIndependentSetOf VertexSet
s EdgeLabeledGraph a
g = [()] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([()] -> Bool) -> [()] -> Bool
forall a b. (a -> b) -> a -> b
$ do
  (Vertex
node1, Vertex
node2, a
_) <- EdgeLabeledGraph a -> [(Vertex, Vertex, a)]
forall a. EdgeLabeledGraph a -> [Edge a]
graphToUnorderedEdges EdgeLabeledGraph a
g
  Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Vertex
node1 Vertex -> VertexSet -> Bool
`IntSet.member` VertexSet
s
  Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Vertex
node2 Vertex -> VertexSet -> Bool
`IntSet.member` VertexSet
s
  () -> [()]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | A clique of a graph is a subset of vertices such that every two distinct vertices in the clique are adjacent.
isCliqueOf :: VertexSet -> EdgeLabeledGraph a -> Bool
isCliqueOf :: forall a. VertexSet -> EdgeLabeledGraph a -> Bool
isCliqueOf VertexSet
s EdgeLabeledGraph a
g = (Vertex -> Bool) -> [Vertex] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Vertex
node -> Vertex -> VertexSet -> VertexSet
IntSet.delete Vertex
node VertexSet
s VertexSet -> VertexSet -> Bool
`IntSet.isSubsetOf` IntMap a -> VertexSet
forall a. IntMap a -> VertexSet
IntMap.keysSet (EdgeLabeledGraph a
g EdgeLabeledGraph a -> Vertex -> IntMap a
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Vertex
node)) (VertexSet -> [Vertex]
IntSet.toList VertexSet
s)