Compressed Sparse Row Graph が最小メモリで使える

競プロ典型 90 問 71 日目 - Fuzzy Priority（★7）

import Control.Monad
import Control.Monad.ST
import qualified Data.ByteString.Char8 as C
import Data.Function
import Data.Foldable
import Data.List
import Data.Sequence (Seq(..), empty, (<|))
import qualified Data.Sequence as Seq
import Data.STRef
import Data.Vector.Unboxed ((!))
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

main = ints >>= \[n,m,k] -> replicateM m ints >>= fmt . sol n m k

ints = unfoldr (C.readInt . C.dropWhile (==' ')) <$> C.getLine

fmt = mapM_ (putStrLn . unwords . map show)

sol n m k abs = runST $ do
  let
    g = mkGraph n . U.fromListN m $ map (\[a,b] -> (a-1,b-1)) abs
  deg <- UM.replicate n (0 :: Int)
  forM_ [0..n-1] $
    U.mapM_ (UM.unsafeModify deg succ) . adj g
  p <- UM.unsafeNew n
  q <- UM.unsafeNew n
  a <- newSTRef empty
  let
    top o j = do
      indeg <- UM.unsafeRead deg j
      if indeg==0 then do
        UM.unsafeWrite q o j
        return $ o+1
      else
        return o

    dn d o i = do
      v <- UM.unsafeRead q i
      U.forM_ (adj g v) $ UM.unsafeModify deg pred
      o' <- pred <$> U.foldM top o (adj g v)
      UM.unsafeWrite p d v
      u <- UM.unsafeRead q o'
      UM.unsafeWrite q i u
      return (o',v,u)

    up i v j u = do
      UM.unsafeWrite q i v
      UM.unsafeWrite q j u
      U.forM_ (adj g v) $ UM.unsafeModify deg succ

    dfs d o
      | d==n      = U.freeze p >>= modifySTRef' a . (<|) >> return True
      | o==0      = return False
      | otherwise = fix (\loop i ->
          if i>=o then
            return True
          else do
            ps <- readSTRef a
            if Seq.length ps>=k then
              return True
            else do
              (o',v,u) <- dn d o i
              b <- dfs (d+1) o'
              if not b then
                return False
              else do
                up i v o' u
                loop (i+1)) 0

  dfs 0 =<< foldM top 0 [0..n-1]
  ps <- readSTRef a
  return $ if k>Seq.length ps then [[-1]] else
    map (U.toList . U.map succ) $ toList ps
    
data Graph = Graph Int (U.Vector Int) (U.Vector Int)

mkGraph :: Int -> U.Vector (Int, Int) -> Graph
mkGraph n xs = runST $ do
  os <- UM.replicate n 0
  U.forM_ xs $ \(u,_) ->
    UM.unsafeModify os succ u
  cs <- U.scanl (+) 0 <$> U.freeze os
  es <- UM.unsafeNew (cs!n)
  U.forM_ xs $ \(u,v) -> do
    UM.unsafeModify os pred u
    o <- UM.unsafeRead os u
    UM.unsafeWrite es (cs!u+o) v
  Graph n cs <$> U.unsafeFreeze es

adj :: Graph -> Int -> U.Vector Int
adj (Graph _ cs es) v = U.unsafeSlice (cs!v) (cs!(v+1)-cs!v) es