SHARE
TWEET

Untitled

a guest Apr 21st, 2019 87 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import Control.Monad
  2. import Control.Monad.Fix
  3. import Control.Monad.Primitive
  4. import Control.Monad.ST
  5. import Data.Primitive.MutVar
  6. import qualified Data.Vector as V
  7. import qualified Data.Vector.Generic.Mutable as VM
  8. import qualified Data.Vector.Unboxed.Mutable as VUM
  9. import qualified Data.Vector.Unboxed as VU
  10. import Data.List
  11.  
  12. newtype LVector s a = LVector (VU.MVector s a, MutVar s Int)
  13.  
  14. new :: (VU.Unbox a, PrimMonad m) => Int -> m (LVector (PrimState m) a)
  15. new n = do
  16.   vec <- VUM.new n
  17.   len <- newMutVar 0
  18.   return $ LVector (vec, len)
  19.  
  20. freeze
  21.   :: (VU.Unbox a, PrimMonad m) => LVector (PrimState m) a -> m (VU.Vector a)
  22. freeze (LVector (v, mv)) = do
  23.   m <- readMutVar mv
  24.   fmap (VU.take m) $ VU.freeze v
  25.  
  26. push :: (VU.Unbox a, PrimMonad m) => LVector (PrimState m) a -> a -> m ()
  27. push lv@(LVector (vec, mv)) a = do
  28.   m <- readMutVar mv
  29.   when (VUM.length vec == m) $ do
  30.     VUM.grow vec (VUM.length vec)
  31.     return ()
  32.  
  33.   VUM.write vec m a
  34.   writeMutVar mv (m + 1)
  35.  
  36. -------
  37.  
  38. type Graph node = [(node,[node])]
  39.  
  40. construct
  41.   :: Eq node => Graph node -> (V.Vector node, GraphInternal, GraphInternal)
  42. construct gnode = runST $ do
  43.   let ns = V.fromList $ map fst gnode
  44.   let rs = V.foldl' (\acc (v, k) i -> if i == k then v else acc i) (\_ -> (-1))
  45.         $ V.indexed ns
  46.  
  47.   g  <- VM.replicateM (V.length ns) (new 10)
  48.   gt <- VM.replicateM (V.length ns) (new 10)
  49.  
  50.   flip mapM_ (zip [0 ..] gnode) $ \(i, (_, es)) -> do
  51.     flip mapM_ (map rs es) $ \e -> do
  52.       gi <- VM.read g i
  53.       push gi e
  54.  
  55.       gti <- VM.read gt e
  56.       push gti i
  57.  
  58.   (,,)
  59.     <$> pure ns
  60.     <*> (V.mapM freeze =<< V.freeze g)
  61.     <*> (V.mapM freeze =<< V.freeze gt)
  62.  
  63.  
  64. type GraphInternal = V.Vector (VU.Vector Int)
  65.  
  66. scc :: Eq node => Graph node -> V.Vector (V.Vector node)
  67. scc gr = V.map (V.map (mapper V.!) . V.convert) build
  68.  where
  69.   (mapper, g, gt) = construct gr
  70.  
  71.   order :: VU.Vector Int
  72.   order = VU.reverse $ runST $ do
  73.     used  <- VUM.replicate (V.length g) False
  74.     order <- new 10
  75.  
  76.     flip fix (0 :: Int) $ \f i -> do
  77.       b <- VM.read used i
  78.       if b
  79.         then return (-1)
  80.         else do
  81.           VM.write used i True
  82.           VU.mapM_ f     (g V.! i)
  83.  
  84.           push     order i
  85.  
  86.           if i == V.length g - 1 then return 0 else f (i + 1)
  87.  
  88.     freeze order
  89.  
  90.   computed :: VU.Vector Int
  91.   size :: Int
  92.   (computed, size) = runST $ do
  93.     comp <- VUM.replicate (V.length g) (-1)
  94.     cnt  <- newMutVar 0
  95.  
  96.     VU.forM_ order $ \o -> do
  97.       b <- VUM.read comp o
  98.       c <- readMutVar cnt
  99.  
  100.       when (b == -1) $ do
  101.         go comp o c
  102.         writeMutVar cnt (c + 1)
  103.  
  104.     (,) <$> VU.freeze comp <*> readMutVar cnt
  105.    where
  106.     go comp idx cnt = do
  107.       b <- VUM.read comp idx
  108.       when (b == -1) $ do
  109.         VUM.write comp idx cnt
  110.         VU.mapM_ (\to -> go comp to cnt) (gt V.! idx)
  111.  
  112.   build :: GraphInternal
  113.   build = runST $ do
  114.     comps <- VM.replicateM size $ new 10
  115.  
  116.     forM_ [0 .. V.length g - 1] $ \i -> do
  117.       VU.forM_ (g V.! i) $ \to -> do
  118.         let x = computed VU.! i
  119.         let y = computed VU.! to
  120.         when (x == y) $ do
  121.           tx <- VM.read comps x
  122.           push tx i
  123.  
  124.     V.mapM freeze =<< V.freeze comps
  125.  
  126. example :: Graph String
  127. example =
  128.   [ ("1", ["2"])
  129.   , ("2", ["3", "5", "6"])
  130.   , ("3", ["4", "7"])
  131.   , ("4", ["3", "8"])
  132.   , ("5", ["1", "6"])
  133.   , ("6", ["7"])
  134.   , ("7", ["6", "8"])
  135.   , ("8", ["8"])
  136.   ]
  137.  
  138. main = print $ scc example
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top