{-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeInType #-} {-# LANGUAGE UnboxedTuples #-} import Control.Monad import Control.Monad.Cont import Control.Monad.Fix import Control.Monad.ST import Data.Bits import Data.Coerce import qualified Data.Ratio as R import GHC.Exts import qualified Data.Vector.Fusion.Stream.Monadic as VFSM import qualified Data.Vector.Generic as VG import qualified Data.Vector.Generic.Mutable as VGM import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as VUM #define MOD 1000000007 modulus :: Num a => a modulus = MOD {-# INLINE modulus #-} main :: IO () main = do [m, k] <- map (read :: String -> Int) . words <$> getLine let as = VU.replicate (m * m) (0 :: Int) matA <- buildSquareMatrix as rep m $ \i -> rep m $ \j -> do VUM.unsafeModify matA succ (i * m + (i + j) `mod` m) VUM.unsafeModify matA succ (i * m + (i * j) `mod` m) matH <- powMat matA (k - 1) mulMat matA matH print =<< VUM.unsafeRead matA 0 type SquareMatrixMint = VUM.IOVector Mint addMat :: SquareMatrixMint -> SquareMatrixMint -> IO () addMat a b = do rep n $ \i -> do item <- VUM.unsafeRead b i VUM.unsafeModify a (+ item) i where !sz = VUM.length a !n = floor . sqrt . fromIntegral $ sz mulMat :: SquareMatrixMint -> SquareMatrixMint -> IO () mulMat a b = do c <- VUM.unsafeNew sz :: IO SquareMatrixMint rep n $ \i -> rep n $ \k -> rep n $ \j -> do item1 <- VUM.unsafeRead a (i * n + k) item2 <- VUM.unsafeRead b (k * n + j) VUM.unsafeModify c (+ item1 * item2) (i * n + j) rep sz $ \i -> VUM.unsafeRead c i >>= VUM.unsafeWrite a i where !sz = VUM.length a !n = floor . sqrt . fromIntegral $ sz powMat :: SquareMatrixMint -> Int -> IO SquareMatrixMint powMat a m = do b <- VUM.replicate sz 0 :: IO SquareMatrixMint rep n $ \i -> VUM.unsafeWrite b (i * n + i) 1 flip fix m $ \loop !i -> do when (i > 0) $ do when (odd i) $ mulMat b a mulMat a a loop (i `div` 2) return b where !sz = VUM.length a !n = floor . sqrt . fromIntegral $ sz buildSquareMatrix :: VU.Vector Int -> IO SquareMatrixMint buildSquareMatrix vec | isSquareMatrix vec = VU.unsafeThaw $ VU.map mint vec | otherwise = VU.unsafeThaw $ VU.map mint vec VU.++ VU.replicate (m * m - n) (0 :: Mint) where !n = VU.length vec !m = succ . floor . sqrt . fromIntegral $ n isSquareMatrix :: VU.Vector Int -> Bool isSquareMatrix vec = let n = VU.length vec in isSquare n isSquare :: Int -> Bool isSquare n = let m = floor . sqrt . fromIntegral $ n in n == m * m infixr 8 ^% infixl 7 *%, /% infixl 6 +%, -% (+%) :: Int -> Int -> Int (I# x#) +% (I# y#) = case x# +# y# of r# -> I# (r# -# ((r# >=# MOD#) *# MOD#)) {-# INLINE (+%) #-} (-%) :: Int -> Int -> Int (I# x#) -% (I# y#) = case x# -# y# of r# -> I# (r# +# ((r# <# 0#) *# MOD#)) {-# INLINE (-%) #-} (*%) :: Int -> Int -> Int (I# x#) *% (I# y#) = case timesWord# (int2Word# x#) (int2Word# y#) of z# -> case timesWord2# z# im# of (# q#, _ #) -> case minusWord# z# (timesWord# q# m#) of v# | isTrue# (geWord# v# m#) -> I# (word2Int# (plusWord# v# m#)) | otherwise -> I# (word2Int# v#) where m# = int2Word# MOD# im# = plusWord# (quotWord# 0xffffffffffffffff## m#) 1## {-# INLINE (*%) #-} (/%) :: Int -> Int -> Int (I# x#) /% (I# y#) = go# y# MOD# 1# 0# where go# a# b# u# v# | isTrue# (b# ># 0#) = case a# `quotInt#` b# of q# -> go# b# (a# -# (q# *# b#)) v# (u# -# (q# *# v#)) | otherwise = I# ((x# *# (u# +# MOD#)) `remInt#` MOD#) {-# INLINE (/%) #-} (^%) :: Int -> Int -> Int x ^% n | n > 0 = go 1 x n | n == 0 = 1 | otherwise = go 1 (1 /% x) (-n) where go !acc !y !m | m .&. 1 == 0 = go acc (y *% y) (unsafeShiftR m 1) | m == 1 = acc *% y | otherwise = go (acc *% y) (y *% y) (unsafeShiftR (m - 1) 1) newtype Mint = Mint { getMint :: Int } deriving newtype (Eq, Ord, Read, Show, Real) mint :: Integral a => a -> Mint mint x = fromIntegral $ mod (fromIntegral x) MOD {-# INLINE mint #-} mintValidate :: Mint -> Bool mintValidate (Mint x) = 0 <= x && x < MOD {-# INLINE mintValidate #-} instance Bounded Mint where minBound = Mint 0 maxBound = Mint $ modulus - 1 instance Enum Mint where toEnum = mint fromEnum = coerce instance Integral Mint where quotRem x y = (x / y, x - x / y * y) toInteger = coerce (toInteger @Int) instance Num Mint where (+) = coerce (+%) (-) = coerce (-%) (*) = coerce (*%) abs = id signum = const (Mint 1) fromInteger x = coerce @Int @Mint . fromInteger $ mod x modulus instance Fractional Mint where (/) = coerce (/%) fromRational q = fromInteger (R.numerator q) / fromInteger (R.denominator q) newtype instance VUM.MVector s Mint = MV_Mint (VUM.MVector s Int) newtype instance VU.Vector Mint = V_Mint (VU.Vector Int) instance VU.Unbox Mint instance VGM.MVector VUM.MVector Mint where basicLength (MV_Mint v) = VGM.basicLength v {-# INLINE basicLength #-} basicUnsafeSlice i n (MV_Mint v) = MV_Mint $ VGM.basicUnsafeSlice i n v {-# INLINE basicUnsafeSlice #-} basicOverlaps (MV_Mint v1) (MV_Mint v2) = VGM.basicOverlaps v1 v2 {-# INLINE basicOverlaps #-} basicUnsafeNew n = MV_Mint `fmap` VGM.basicUnsafeNew n {-# INLINE basicUnsafeNew #-} basicInitialize (MV_Mint v) = VGM.basicInitialize v {-# INLINE basicInitialize #-} basicUnsafeReplicate n x = MV_Mint `fmap` VGM.basicUnsafeReplicate n (coerce x) {-# INLINE basicUnsafeReplicate #-} basicUnsafeRead (MV_Mint v) i = coerce `fmap` VGM.basicUnsafeRead v i {-# INLINE basicUnsafeRead #-} basicUnsafeWrite (MV_Mint v) i x = VGM.basicUnsafeWrite v i (coerce x) {-# INLINE basicUnsafeWrite #-} basicClear (MV_Mint v) = VGM.basicClear v {-# INLINE basicClear #-} basicSet (MV_Mint v) x = VGM.basicSet v (coerce x) {-# INLINE basicSet #-} basicUnsafeCopy (MV_Mint v1) (MV_Mint v2) = VGM.basicUnsafeCopy v1 v2 {-# INLINE basicUnsafeCopy #-} basicUnsafeMove (MV_Mint v1) (MV_Mint v2) = VGM.basicUnsafeMove v1 v2 {-# INLINE basicUnsafeMove #-} basicUnsafeGrow (MV_Mint v) n = MV_Mint `fmap` VGM.basicUnsafeGrow v n {-# INLINE basicUnsafeGrow #-} instance VG.Vector VU.Vector Mint where basicUnsafeFreeze (MV_Mint v) = V_Mint `fmap` VG.basicUnsafeFreeze v {-# INLINE basicUnsafeFreeze #-} basicUnsafeThaw (V_Mint v) = MV_Mint `fmap` VG.basicUnsafeThaw v {-# INLINE basicUnsafeThaw #-} basicLength (V_Mint v) = VG.basicLength v {-# INLINE basicLength #-} basicUnsafeSlice i n (V_Mint v) = V_Mint $ VG.basicUnsafeSlice i n v {-# INLINE basicUnsafeSlice #-} basicUnsafeIndexM (V_Mint v) i = coerce `fmap` VG.basicUnsafeIndexM v i {-# INLINE basicUnsafeIndexM #-} basicUnsafeCopy (MV_Mint mv) (V_Mint v) = VG.basicUnsafeCopy mv v elemseq _ = seq {-# INLINE elemseq #-} rep :: Monad m => Int -> (Int -> m ()) -> m () rep n = flip VFSM.mapM_ (streamG 0 (n - 1) const 0 (+) 1) {-# INLINE rep #-} rep' :: Monad m => Int -> (Int -> m ()) -> m () rep' n = flip VFSM.mapM_ (streamG 0 n const 0 (+) 1) {-# INLINE rep' #-} rep1 :: Monad m => Int -> (Int -> m ()) -> m () rep1 n = flip VFSM.mapM_ (streamG 1 (n - 1) const 0 (+) 1) {-# INLINE rep1 #-} rep1' :: Monad m => Int -> (Int -> m ()) -> m () rep1' n = flip VFSM.mapM_ (streamG 1 n const 0 (+) 1) {-# INLINE rep1' #-} rev :: Monad m => Int -> (Int -> m ()) -> m () rev n = flip VFSM.mapM_ (streamRG (n - 1) 0 const 0 (-) 1) {-# INLINE rev #-} rev' :: Monad m => Int -> (Int -> m ()) -> m () rev' n = flip VFSM.mapM_ (streamRG n 0 const 0 (-) 1) {-# INLINE rev' #-} rev1 :: Monad m => Int -> (Int -> m ()) -> m () rev1 n = flip VFSM.mapM_ (streamRG (n - 1) 1 const 0 (-) 1) {-# INLINE rev1 #-} rev1' :: Monad m => Int -> (Int -> m ()) -> m () rev1' n = flip VFSM.mapM_ (streamRG n 1 const 0 (-) 1) {-# INLINE rev1' #-} range :: Monad m => Int -> Int -> (Int -> m ()) -> m () range l r = flip VFSM.mapM_ (streamG l r const 0 (+) 1) {-# INLINE range #-} rangeR :: Monad m => Int -> Int -> (Int -> m ()) -> m () rangeR r l = flip VFSM.mapM_ (streamRG r l const 0 (-) 1) {-# INLINE rangeR #-} forP :: Monad m => Int -> (Int -> m ()) -> m () forP p = flip VFSM.mapM_ (streamG 2 p (^) 2 (+) 1) {-# INLINE forP #-} forG :: Monad m => Int -> Int -> (Int -> Int -> Int) -> Int -> (Int -> Int -> Int) -> Int -> (Int -> m ()) -> m () forG l r f p g d = flip VFSM.mapM_ (streamG l r f p g d) {-# INLINE forG #-} forRG :: Monad m => Int -> Int -> (Int -> Int -> Int) -> Int -> (Int -> Int -> Int) -> Int -> (Int -> m ()) -> m () forRG r l f p g d = flip VFSM.mapM_ (streamRG r l f p g d) {-# INLINE forRG #-} streamG :: Monad m => Int -> Int -> (Int -> Int -> Int) -> Int -> (Int -> Int -> Int) -> Int -> VFSM.Stream m Int streamG !l !r !f !p !g !d = VFSM.Stream step l where step x | f x p <= r = return $ VFSM.Yield x (g x d) | otherwise = return VFSM.Done {-# INLINE [0] step #-} {-# INLINE [1] streamG #-} streamRG :: Monad m => Int -> Int -> (Int -> Int -> Int) -> Int -> (Int -> Int -> Int) -> Int -> VFSM.Stream m Int streamRG !r !l !f !p !g !d = VFSM.Stream step r where step x | f x p >= l = return $ VFSM.Yield x (g x d) | otherwise = return VFSM.Done {-# INLINE [0] step #-} {-# INLINE [1] streamRG #-} withBreakIO :: ((r -> ContT r IO b) -> ContT r IO r) -> IO r withBreakIO = flip runContT pure . callCC {-# INLINE withBreakIO #-} withBreakST :: ((r -> ContT r (ST s) b) -> ContT r (ST s) r) -> (ST s) r withBreakST = flip runContT pure . callCC {-# INLINE withBreakST #-}