{-# 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
  ( EdgeLabeledGraph
  , Graph
  , graphToUnorderedEdges
  , graphFromUnorderedEdges
  , graphFromUnorderedEdgesWith
  , isIndependentSet
  ) 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)

type EdgeLabeledGraph a = Array Int (IntMap a)

type Graph = EdgeLabeledGraph ()

graphToUnorderedEdges :: EdgeLabeledGraph a -> [(Int, Int, a)]
graphToUnorderedEdges :: EdgeLabeledGraph a -> [(Int, Int, a)]
graphToUnorderedEdges EdgeLabeledGraph a
g = do
  (Int
node1, IntMap a
nodes) <- EdgeLabeledGraph a -> [(Int, IntMap a)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs EdgeLabeledGraph a
g
  (Int
node2, a
a) <- IntMap a -> [(Int, a)]
forall a. IntMap a -> [(Int, a)]
IntMap.toList (IntMap a -> [(Int, a)]) -> IntMap a -> [(Int, a)]
forall a b. (a -> b) -> a -> b
$ (IntMap a, IntMap a) -> IntMap a
forall a b. (a, b) -> b
snd ((IntMap a, IntMap a) -> IntMap a)
-> (IntMap a, IntMap a) -> IntMap a
forall a b. (a -> b) -> a -> b
$ Int -> IntMap a -> (IntMap a, IntMap a)
forall a. Int -> IntMap a -> (IntMap a, IntMap a)
IntMap.split Int
node1 IntMap a
nodes
  (Int, Int, a) -> [(Int, Int, a)]
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
node1, Int
node2, a
a)

graphFromUnorderedEdges :: Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdges :: Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdges = (a -> a -> a) -> Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
forall a.
(a -> a -> a) -> Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith a -> a -> a
forall a b. a -> b -> a
const

graphFromUnorderedEdgesWith :: (a -> a -> a) -> Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith :: (a -> a -> a) -> Int -> [(Int, Int, a)] -> EdgeLabeledGraph a
graphFromUnorderedEdgesWith a -> a -> a
f Int
n [(Int, Int, a)]
es = (forall s. ST s (STArray s Int (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 Int (IntMap a))) -> EdgeLabeledGraph a)
-> (forall s. ST s (STArray s Int (IntMap a)))
-> EdgeLabeledGraph a
forall a b. (a -> b) -> a -> b
$ do
  STArray s Int (IntMap a)
a <- (Int, Int) -> IntMap a -> ST s (STArray s Int (IntMap a))
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) IntMap a
forall a. IntMap a
IntMap.empty
  let ins :: Int -> Int -> a -> m ()
ins Int
i Int
x a
l = do
        IntMap a
m <- STArray s Int (IntMap a) -> Int -> m (IntMap a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Int (IntMap a)
a Int
i
        STArray s Int (IntMap a) -> Int -> IntMap a -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Int (IntMap a)
a Int
i (IntMap a -> m ()) -> IntMap a -> m ()
forall a b. (a -> b) -> a -> b
$! (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith a -> a -> a
f Int
x a
l IntMap a
m
  [(Int, Int, a)] -> ((Int, Int, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Int, Int, a)]
es (((Int, Int, a) -> ST s ()) -> ST s ())
-> ((Int, Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Int
node1, Int
node2, a
a) -> do
    Int -> Int -> a -> ST s ()
forall (m :: * -> *).
MArray (STArray s) (IntMap a) m =>
Int -> Int -> a -> m ()
ins Int
node1 Int
node2 a
a
    Int -> Int -> a -> ST s ()
forall (m :: * -> *).
MArray (STArray s) (IntMap a) m =>
Int -> Int -> a -> m ()
ins Int
node2 Int
node1 a
a
  STArray s Int (IntMap a) -> ST s (STArray s Int (IntMap a))
forall (m :: * -> *) a. Monad m => a -> m a
return STArray s Int (IntMap a)
a

isIndependentSet :: EdgeLabeledGraph a -> IntSet -> Bool
isIndependentSet :: EdgeLabeledGraph a -> IntSet -> Bool
isIndependentSet EdgeLabeledGraph a
g IntSet
s = [()] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([()] -> Bool) -> [()] -> Bool
forall a b. (a -> b) -> a -> b
$ do
  (Int
node1, Int
node2, a
_) <- EdgeLabeledGraph a -> [(Int, Int, a)]
forall a. EdgeLabeledGraph a -> [(Int, Int, a)]
graphToUnorderedEdges EdgeLabeledGraph a
g
  Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Int
node1 Int -> IntSet -> Bool
`IntSet.member` IntSet
s
  Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> [()]) -> Bool -> [()]
forall a b. (a -> b) -> a -> b
$ Int
node2 Int -> IntSet -> Bool
`IntSet.member` IntSet
s
  () -> [()]
forall (m :: * -> *) a. Monad m => a -> m a
return ()