Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 260b53c

Browse files
committed
Add TDPC-T
1 parent 54fcd03 commit 260b53c

File tree

2 files changed

+227
-0
lines changed

2 files changed

+227
-0
lines changed

‎tdpc-t/Main.hs‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
-- https://github.com/minoki/my-atcoder-solutions
2+
{-# LANGUAGE BangPatterns #-}
3+
{-# LANGUAGE DataKinds #-}
4+
{-# LANGUAGE NoStarIsType #-}
5+
{-# LANGUAGE TypeFamilies #-}
6+
{-# LANGUAGE TypeOperators #-}
7+
import qualified Data.ByteString.Char8 as BS
8+
import Data.Char (isSpace)
9+
import Data.Int (Int64)
10+
import Data.List (tails, unfoldr)
11+
import qualified Data.Vector.Unboxing as U
12+
import GHC.TypeNats (type (+), KnownNat, Nat, type (^),
13+
natVal)
14+
15+
type Poly = U.Vector (IntMod (10^9 + 7))
16+
17+
-- 多項式は
18+
-- U.fromList [a,b,c,...,z] = a + b * X + c * X^2 + ... + z * X^(k-1)
19+
-- により表す。
20+
21+
reduce :: Int -> Poly -> Poly
22+
reduce !k !v | U.last v == 0 = U.init v
23+
| U.length v <= k = v
24+
| otherwise = let b = U.last v
25+
l = U.length v
26+
in reduce k (U.imap (\i a -> if i >= l - k - 1 then a + b else a) (U.init v))
27+
28+
-- 多項式の積を X^k - X^(k-1) - ... - X - 1 で割った余りを返す。
29+
mulP :: Int -> Poly -> Poly -> Poly
30+
mulP !k !v !w = reduce k $ U.generate (U.length v + U.length w - 1) $
31+
\i -> sum [(v U.! (i-j)) * (w U.! j) | j <- [0..U.length w-1], j <= i, j > i - U.length v]
32+
33+
-- 多項式に X をかけたものを X^k - X^(k-1) - ... - X - 1 で割った余りを返す。
34+
mulByX :: Int -> Poly -> Poly
35+
mulByX !k !v
36+
| U.length v == k = let !v_k = v U.! (k-1)
37+
in U.generate k $ \i -> if i == 0 then
38+
v_k
39+
else
40+
v_k + (v U.! (i - 1))
41+
| otherwise = U.cons 0 v
42+
43+
-- X の(mod X^k - X^(k-1) - ... - X - 1 での)n 乗
44+
powX :: Int -> Int -> Poly
45+
powX !k !n = doPowX n
46+
where
47+
doPowX 0 = U.fromList [1] -- 1
48+
doPowX 1 = U.fromList [0,1] -- X
49+
doPowX i = case i `quotRem` 2 of
50+
(j,0) -> let !f = doPowX j -- X^j mod P
51+
in mulP k f f
52+
(j,_) -> let !f = doPowX j -- X^j mod P
53+
in mulByX k (mulP k f f)
54+
55+
main :: IO ()
56+
main = do
57+
[k,n] <- unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine
58+
if n <= k then
59+
print 1
60+
else do
61+
let f = powX k (n - k) -- X^(n-k) mod X^k - X^(k-1) - ... - X - 1
62+
let seq = replicate k 1 ++ map (sum . take k) (tails seq) -- 数列
63+
print $ sum $ zipWith (*) (U.toList f) (drop (k-1) seq)
64+
65+
--
66+
-- Modular Arithmetic
67+
--
68+
69+
newtype IntMod (m :: Nat) = IntMod { unwrapN :: Int64 } deriving (Eq)
70+
71+
instance Show (IntMod m) where
72+
show (IntMod x) = show x
73+
74+
instance KnownNat m => Num (IntMod m) where
75+
t@(IntMod x) + IntMod y
76+
| x + y >= modulus = IntMod (x + y - modulus)
77+
| otherwise = IntMod (x + y)
78+
where modulus = fromIntegral (natVal t)
79+
t@(IntMod x) - IntMod y
80+
| x >= y = IntMod (x - y)
81+
| otherwise = IntMod (x - y + modulus)
82+
where modulus = fromIntegral (natVal t)
83+
t@(IntMod x) * IntMod y = IntMod ((x * y) `rem` modulus)
84+
where modulus = fromIntegral (natVal t)
85+
fromInteger n = let result = IntMod (fromInteger (n `mod` fromIntegral modulus))
86+
modulus = natVal result
87+
in result
88+
abs = undefined; signum = undefined
89+
{-# SPECIALIZE instance Num (IntMod 1000000007) #-}
90+
91+
fromIntegral_Int64_IntMod :: KnownNat m => Int64 -> IntMod m
92+
fromIntegral_Int64_IntMod n = result
93+
where
94+
result | 0 <= n && n < modulus = IntMod n
95+
| otherwise = IntMod (n `mod` modulus)
96+
modulus = fromIntegral (natVal result)
97+
98+
{-# RULES
99+
"fromIntegral/Int->IntMod" fromIntegral = fromIntegral_Int64_IntMod . (fromIntegral :: Int -> Int64) :: Int -> IntMod (10^9 + 7)
100+
"fromIntegral/Int64->IntMod" fromIntegral = fromIntegral_Int64_IntMod :: Int64 -> IntMod (10^9 + 7)
101+
#-}
102+
103+
instance U.Unboxable (IntMod m) where
104+
type Rep (IntMod m) = Int64

‎tdpc-t/Mutable.hs‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
-- https://github.com/minoki/my-atcoder-solutions
2+
{-# LANGUAGE BangPatterns #-}
3+
{-# LANGUAGE DataKinds #-}
4+
{-# LANGUAGE NoStarIsType #-}
5+
{-# LANGUAGE TypeFamilies #-}
6+
{-# LANGUAGE TypeOperators #-}
7+
import Control.Monad
8+
import Control.Monad.ST
9+
import qualified Data.ByteString.Char8 as BS
10+
import Data.Char (isSpace)
11+
import Data.Coerce
12+
import Data.Int (Int64)
13+
import Data.List (foldl', tails, unfoldr)
14+
import qualified Data.Vector.Unboxing as U
15+
import qualified Data.Vector.Unboxing.Mutable as UM
16+
import GHC.TypeNats (type (+), KnownNat, Nat,
17+
type (^), natVal)
18+
19+
type Poly = U.Vector (IntMod (10^9 + 7))
20+
type PolyM s = UM.MVector s (IntMod (10^9 + 7))
21+
22+
sum' :: KnownNat m => [IntMod m] -> IntMod m
23+
sum' = fromIntegral . foldl' (\x y -> x + unwrapN y) 0
24+
{-# INLINE sum' #-}
25+
26+
-- 多項式は
27+
-- U.fromList [a,b,c,...,z] = a + b * X + c * X^2 + ... + z * X^(k-1)
28+
-- により表す。
29+
30+
-- 多項式を X^k - X^(k-1) - ... - X - 1 で割った余りを返す。
31+
reduceM :: Int -> PolyM s -> ST s (PolyM s)
32+
reduceM !k !v = loop (UM.length v)
33+
where loop !l | l <= k = return (UM.take l v)
34+
| otherwise = do b <- UM.read v (l - 1)
35+
forM_ [l - k - 1 .. l - 2] $ \i -> do
36+
UM.modify v (+ b) i
37+
loop (l - 1)
38+
39+
-- 多項式の積を X^k - X^(k-1) - ... - X - 1 で割った余りを返す。
40+
mulP :: Int -> Poly -> Poly -> Poly
41+
mulP !k !v !w = {- U.force $ -} U.create $ do
42+
let !vl = U.length v
43+
!wl = U.length w
44+
s <- UM.new (vl + wl - 1)
45+
forM_ [0 .. vl + wl - 2] $ \i -> do
46+
let !x = sum' [(v U.! (i-j)) * (w U.! j) | j <- [max 0 (i - vl + 1) .. min (wl - 1) i]]
47+
UM.write s i x
48+
reduceM k s
49+
50+
-- 多項式に X をかけたものを X^k - X^(k-1) - ... - X - 1 で割った余りを返す。
51+
mulByX :: Int -> Poly -> Poly
52+
mulByX !k !v
53+
| U.length v == k = let !v_k = v U.! (k-1)
54+
in U.generate k $ \i -> if i == 0 then
55+
v_k
56+
else
57+
v_k + (v U.! (i - 1))
58+
| otherwise = U.cons 0 v
59+
60+
-- X の(mod X^k - X^(k-1) - ... - X - 1 での)n 乗
61+
powX :: Int -> Int -> Poly
62+
powX !k !n = doPowX n
63+
where
64+
doPowX 0 = U.fromList [1] -- 1
65+
doPowX 1 = U.fromList [0,1] -- X
66+
doPowX i = case i `quotRem` 2 of
67+
(j,0) -> let !f = doPowX j -- X^j mod P
68+
in mulP k f f
69+
(j,_) -> let !f = doPowX j -- X^j mod P
70+
in mulByX k (mulP k f f)
71+
72+
main :: IO ()
73+
main = do
74+
[k,n] <- unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine
75+
-- 2 <= k <= 1000
76+
-- 1 <= n <= 10^9
77+
if n <= k then
78+
print 1
79+
else do
80+
let f = powX k (n - k) -- X^(n-k) mod X^k - X^(k-1) - ... - X - 1
81+
let seq = replicate k 1 ++ map (sum . take k) (tails seq) -- 数列
82+
print $ sum $ zipWith (*) (U.toList f) (drop (k-1) seq)
83+
84+
--
85+
-- Modular Arithmetic
86+
--
87+
88+
newtype IntMod (m :: Nat) = IntMod { unwrapN :: Int64 } deriving (Eq)
89+
90+
instance Show (IntMod m) where
91+
show (IntMod x) = show x
92+
93+
instance KnownNat m => Num (IntMod m) where
94+
t@(IntMod x) + IntMod y
95+
| x + y >= modulus = IntMod (x + y - modulus)
96+
| otherwise = IntMod (x + y)
97+
where modulus = fromIntegral (natVal t)
98+
t@(IntMod x) - IntMod y
99+
| x >= y = IntMod (x - y)
100+
| otherwise = IntMod (x - y + modulus)
101+
where modulus = fromIntegral (natVal t)
102+
t@(IntMod x) * IntMod y = IntMod ((x * y) `rem` modulus)
103+
where modulus = fromIntegral (natVal t)
104+
fromInteger n = let result = IntMod (fromInteger (n `mod` fromIntegral modulus))
105+
modulus = natVal result
106+
in result
107+
abs = undefined; signum = undefined
108+
{-# SPECIALIZE instance Num (IntMod 1000000007) #-}
109+
110+
fromIntegral_Int64_IntMod :: KnownNat m => Int64 -> IntMod m
111+
fromIntegral_Int64_IntMod n = result
112+
where
113+
result | 0 <= n && n < modulus = IntMod n
114+
| otherwise = IntMod (n `mod` modulus)
115+
modulus = fromIntegral (natVal result)
116+
117+
{-# RULES
118+
"fromIntegral/Int->IntMod" fromIntegral = fromIntegral_Int64_IntMod . (fromIntegral :: Int -> Int64) :: Int -> IntMod (10^9 + 7)
119+
"fromIntegral/Int64->IntMod" fromIntegral = fromIntegral_Int64_IntMod :: Int64 -> IntMod (10^9 + 7)
120+
#-}
121+
122+
instance U.Unboxable (IntMod m) where
123+
type Rep (IntMod m) = Int64

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /