Last active
March 9, 2017 04:01
-
-
Save AndrasKovacs/582808b6b5cc67bc36a2 to your computer and use it in GitHub Desktop.
Graph building code for SO question. See link in source.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
-- http://stackoverflow.com/questions/24278006/need-advice-on-optimising-haskell-data-processing/ | |
-- http://stackoverflow.com/questions/24344440/optimising-haskell-data-processing-part-ii | |
-- ****************** "Normal" version ***************************** | |
import Data.Vector (Vector) | |
import qualified Data.Vector as V | |
import qualified Data.ByteString.Char8 as BS | |
import qualified Data.IntSet as IS | |
import Control.Monad.State | |
import Control.Monad | |
import Control.Arrow | |
data Node = Node {bwd, fwd :: Edges} deriving (Eq, Show) | |
type Edges = [Int] | |
type Graph = Vector Node | |
type Stack = [Int] | |
main = do | |
edges <- V.fromList `fmap` getEdges "SCC.txt" | |
let maxIndex = fst $ V.last edges | |
let graph = createGraph maxIndex edges | |
putStrLn "Number of components in graph:" | |
print $ length $ kosaraju graph | |
getEdges :: String -> IO [(Int, Int)] | |
getEdges path = do | |
lines <- (map BS.words . BS.lines) `fmap` BS.readFile path | |
let pairs = (map . map) (maybe (error "can't read Int") fst . BS.readInt) lines | |
return [(a, b) | [a, b] <- pairs] | |
createGraph :: Int -> Vector (Int, Int) -> Graph | |
createGraph maxIndex edges = graph'' where | |
graph = V.replicate (maxIndex + 1) (Node [] []) | |
graph' = V.accumulate (\(Node f b) x -> Node (x:f) b) graph edges | |
graph'' = V.accumulate (\(Node f b) x -> Node f (x:b)) graph' (V.map (\(a, b) -> (b, a)) edges) | |
-- Step 1 : push the nodes onto a stack in order of dfs completion. | |
dfsReorder :: Graph -> Stack | |
dfsReorder graph = snd $ execState (mapM_ go [0 .. V.length graph - 1]) (IS.empty, []) where | |
go :: Int -> State (IS.IntSet, Stack) () | |
go i = do | |
visited <- gets (IS.member i . fst) | |
unless visited $ do | |
modify $ first $ IS.insert i | |
mapM_ go $ fwd $ graph V.! i | |
modify $ second (i:) | |
-- Step 2 : do dfs on the stack, but using a transposed graph. | |
-- here we return a list of representative nodes for the graph components. | |
-- we use filterM to skip over the visited nodes in the stack. | |
kosaraju :: Graph -> [Int] | |
kosaraju graph = evalState (filterM go (dfsReorder graph)) IS.empty where | |
go :: Int -> State IS.IntSet Bool | |
go i = do | |
visited <- gets $ IS.member i | |
if visited then | |
return False | |
else do | |
modify $ IS.insert i | |
mapM_ go $ bwd $ graph V.! i | |
return True | |
-- ****************** Optimized/mutable version ********************* | |
-- Works the same as the "normal" version, except here we have mutable | |
-- state instead of the State monad. I also sprinkled in | |
-- unsafe read/write ops for maximum performance. It's about 5 times faster | |
-- than the other version. | |
{-- | |
import Data.Vector (Vector) | |
import qualified Data.Vector as V | |
import qualified Data.Vector.Unboxed as UV | |
import qualified Data.Vector.Unboxed.Mutable as MUV | |
import qualified Data.ByteString.Char8 as BS | |
import Control.Monad | |
import Control.Monad.ST | |
import Data.STRef | |
import Data.Function | |
type Edges = UV.Vector (Int, Int) | |
type Stack = [Int] | |
type Graph = Vector [Int] | |
main = do | |
putStrLn "Number of strongly connected components in graph:" | |
print . length . kosaraju =<< getEdges "SCC.txt" | |
getEdges :: String -> IO Edges | |
getEdges path = do | |
lines <- (map BS.words . BS.lines) `fmap` BS.readFile path | |
let pairs = (map . map) (maybe (error "can't read Int") fst . BS.readInt) lines | |
return $ UV.fromList $ map (\[a, b] -> (a, b)) pairs | |
createGraph :: Edges -> Graph | |
createGraph edges = | |
V.unsafeAccumulate (flip (:)) | |
(V.replicate (uncurry max (UV.last edges) + 1) []) | |
(V.convert edges) -- conversion from unboxed to boxed vector | |
dfsReorder :: Graph -> Stack | |
dfsReorder graph = runST $ do | |
let nodeNum = V.length graph | |
stack <- newSTRef [] | |
visited <- MUV.unsafeNew nodeNum :: ST s (MUV.STVector s Bool) | |
MUV.set visited False | |
forM_ [0 .. nodeNum - 1] $ fix $ \go i -> do | |
isVisited <- MUV.unsafeRead visited i | |
unless isVisited $ do | |
MUV.unsafeWrite visited i True | |
mapM_ go $ graph `V.unsafeIndex` i | |
modifySTRef' stack (i:) | |
readSTRef stack | |
kosaraju :: Edges -> [Int] | |
kosaraju edges = runST $ do | |
let graph = createGraph edges | |
graph' = createGraph $ UV.map (\(a, b) -> (b, a)) edges | |
visited <- MUV.unsafeNew (V.length graph) :: ST s (MUV.STVector s Bool) | |
MUV.set visited False | |
flip filterM (dfsReorder graph) $ fix $ \go i -> do | |
isVisited <- MUV.unsafeRead visited i | |
if isVisited then | |
return False | |
else do | |
MUV.unsafeWrite visited i True | |
mapM_ go $ graph' `V.unsafeIndex` i | |
return True | |
--} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment