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 688ea9e

Browse files
committed
AtCoder Typical Contest 001 - C - 高速フーリエ変換
1 parent bf187ef commit 688ea9e

File tree

3 files changed

+412
-0
lines changed

3 files changed

+412
-0
lines changed

‎README.md‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -369,3 +369,11 @@ Haskellで競技プログラミングをやるテクニックは「[Haskellで
369369
* [ ] J - Segment Tree
370370
* [ ] K - Range Affine Range Sum
371371
* [ ] L - Lazy Segment Tree
372+
373+
## AtCoder Typical Contest 001
374+
375+
<https://atcoder.jp/contests/atc001>
376+
377+
* [ ] A - 深さ優先探索
378+
* [ ] B - Union Find
379+
* [x] C - 高速フーリエ変換

‎atc001-c/Karatsuba.hs‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
-- https://github.com/minoki/my-atcoder-solutions
2+
{-# LANGUAGE BangPatterns #-}
3+
{-# LANGUAGE MultiParamTypeClasses #-}
4+
{-# LANGUAGE ScopedTypeVariables #-}
5+
{-# LANGUAGE TypeApplications #-}
6+
import Control.Monad
7+
import qualified Data.ByteString.Char8 as BS
8+
import Data.Char (isSpace)
9+
import Data.Coerce
10+
import Data.Int (Int64)
11+
import Data.List (unfoldr)
12+
import qualified Data.Vector.Generic as G
13+
import qualified Data.Vector.Unboxed as U
14+
import qualified Data.Vector.Unboxed.Mutable as UM
15+
16+
main = do
17+
n <- readLn @Int -- n <= 10^5
18+
(as,bs) <- fmap U.unzip $ U.replicateM n $ do
19+
[a,b] <- unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine
20+
return (a,b)
21+
let p, q :: Poly U.Vector Int
22+
p = Poly $ normalizePoly (0 `U.cons` as)
23+
q = Poly $ normalizePoly (0 `U.cons` bs)
24+
!v = coeffAsc (p * q)
25+
!l = U.length v
26+
forM_ [1..2*n] $ \k -> do
27+
print $ if k < l then
28+
v U.! k
29+
else
30+
0
31+
32+
--
33+
-- Univariate polynomial
34+
--
35+
36+
newtype Poly vec a = Poly { coeffAsc :: vec a } deriving Eq
37+
38+
normalizePoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a
39+
normalizePoly v | G.null v || G.last v /= 0 = v
40+
| otherwise = normalizePoly (G.init v)
41+
42+
addPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
43+
addPoly v w = case compare n m of
44+
LT -> G.generate m $ \i -> if i < n
45+
then v G.! i + w G.! i
46+
else w G.! i
47+
GT -> G.generate n $ \i -> if i < m
48+
then v G.! i + w G.! i
49+
else v G.! i
50+
EQ -> normalizePoly $ G.zipWith (+) v w
51+
where n = G.length v
52+
m = G.length w
53+
54+
subPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
55+
subPoly v w = case compare n m of
56+
LT -> G.generate m $ \i -> if i < n
57+
then v G.! i - w G.! i
58+
else negate (w G.! i)
59+
GT -> G.generate n $ \i -> if i < m
60+
then v G.! i - w G.! i
61+
else v G.! i
62+
EQ -> normalizePoly $ G.zipWith (-) v w
63+
where n = G.length v
64+
m = G.length w
65+
66+
naiveMulPoly :: (Num a, G.Vector vec a) => vec a -> vec a -> vec a
67+
naiveMulPoly v w = G.generate (n + m - 1) $
68+
\i -> sum [(v G.! (i-j)) * (w G.! j) | j <- [max (i-n+1) 0..min i (m-1)]]
69+
where n = G.length v
70+
m = G.length w
71+
72+
doMulP :: (Eq a, Num a, G.Vector vec a) => Int -> vec a -> vec a -> vec a
73+
doMulP n !v !w | n <= 16 = naiveMulPoly v w
74+
doMulP n !v !w
75+
| G.null v = v
76+
| G.null w = w
77+
| G.length v < n2 = let (w0, w1) = G.splitAt n2 w
78+
u0 = doMulP n2 v w0
79+
u1 = doMulP n2 v w1
80+
in G.generate (G.length v + G.length w - 1)
81+
$ \i -> case () of
82+
_ | i < n2 -> u0 `at` i
83+
| i < n -> (u0 `at` i) + (u1 `at` (i - n2))
84+
| i < n + n2 -> (u1 `at` (i - n2))
85+
| G.length w < n2 = let (v0, v1) = G.splitAt n2 v
86+
u0 = doMulP n2 v0 w
87+
u1 = doMulP n2 v1 w
88+
in G.generate (G.length v + G.length w - 1)
89+
$ \i -> case () of
90+
_ | i < n2 -> u0 `at` i
91+
| i < n -> (u0 `at` i) + (u1 `at` (i - n2))
92+
| i < n + n2 -> (u1 `at` (i - n2))
93+
| otherwise = let (v0, v1) = G.splitAt n2 v
94+
(w0, w1) = G.splitAt n2 w
95+
v0_1 = v0 `addPoly` v1
96+
w0_1 = w0 `addPoly` w1
97+
p = doMulP n2 v0_1 w0_1
98+
q = doMulP n2 v0 w0
99+
r = doMulP n2 v1 w1
100+
-- s = (p `subPoly` q) `subPoly` r -- p - q - r
101+
-- q + s*X^n2 + r*X^n
102+
in G.generate (G.length v + G.length w - 1)
103+
$ \i -> case () of
104+
_ | i < n2 -> q `at` i
105+
| i < n -> ((q `at` i) + (p `at` (i - n2))) - ((q `at` (i - n2)) + (r `at` (i - n2)))
106+
| i < n + n2 -> ((r `at` (i - n)) + (p `at` (i - n2))) - ((q `at` (i - n2)) + (r `at` (i - n2)))
107+
| otherwise -> r `at` (i - n)
108+
where n2 = n `quot` 2
109+
at :: (Num a, G.Vector vec a) => vec a -> Int -> a
110+
at v i = if i < G.length v then v G.! i else 0
111+
{-# INLINE doMulP #-}
112+
113+
mulPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
114+
mulPoly !v !w = let k = ceiling ((log (fromIntegral (max n m)) :: Double) / log 2) :: Int
115+
in doMulP (2^k) v w
116+
where n = G.length v
117+
m = G.length w
118+
{-# INLINE mulPoly #-}
119+
120+
zeroPoly :: (G.Vector vec a) => Poly vec a
121+
zeroPoly = Poly G.empty
122+
123+
constPoly :: (Eq a, Num a, G.Vector vec a) => a -> Poly vec a
124+
constPoly 0 = Poly G.empty
125+
constPoly x = Poly (G.singleton x)
126+
127+
scalePoly :: (Eq a, Num a, G.Vector vec a) => a -> Poly vec a -> Poly vec a
128+
scalePoly a (Poly xs)
129+
| a == 0 = zeroPoly
130+
| otherwise = Poly $ G.map (* a) xs
131+
132+
valueAtPoly :: (Num a, G.Vector vec a) => Poly vec a -> a -> a
133+
valueAtPoly (Poly xs) t = G.foldr' (\a b -> a + t * b) 0 xs
134+
135+
instance (Eq a, Num a, G.Vector vec a) => Num (Poly vec a) where
136+
(+) = coerce (addPoly :: vec a -> vec a -> vec a)
137+
(-) = coerce (subPoly :: vec a -> vec a -> vec a)
138+
negate (Poly v) = Poly (G.map negate v)
139+
(*) = coerce (mulPoly :: vec a -> vec a -> vec a)
140+
fromInteger = constPoly . fromInteger
141+
abs = undefined; signum = undefined
142+
143+
divModPoly :: (Eq a, Fractional a, G.Vector vec a) => Poly vec a -> Poly vec a -> (Poly vec a, Poly vec a)
144+
divModPoly f g@(Poly w)
145+
| G.null w = error "divModPoly: divide by zero"
146+
| degree f < degree g = (zeroPoly, f)
147+
| otherwise = loop zeroPoly (scalePoly (recip b) f)
148+
where
149+
g' = toMonic g
150+
b = leadingCoefficient g
151+
-- invariant: f == q * g + scalePoly b p
152+
loop q p | degree p < degree g = (q, scalePoly b p)
153+
| otherwise = let q' = Poly (G.drop (degree' g) (coeffAsc p))
154+
in loop (q + q') (p - q' * g')
155+
156+
toMonic :: (Fractional a, G.Vector vec a) => Poly vec a -> Poly vec a
157+
toMonic f@(Poly xs)
158+
| G.null xs = zeroPoly
159+
| otherwise = Poly $ G.map (* recip (leadingCoefficient f)) xs
160+
161+
leadingCoefficient :: (Num a, G.Vector vec a) => Poly vec a -> a
162+
leadingCoefficient (Poly xs)
163+
| G.null xs = 0
164+
| otherwise = G.last xs
165+
166+
degree :: G.Vector vec a => Poly vec a -> Maybe Int
167+
degree (Poly xs) = case G.length xs - 1 of
168+
-1 -> Nothing
169+
n -> Just n
170+
171+
degree' :: G.Vector vec a => Poly vec a -> Int
172+
degree' (Poly xs) = case G.length xs of
173+
0 -> error "degree': zero polynomial"
174+
n -> n - 1
175+
176+
-- 組立除法
177+
-- second constPoly (divModByDeg1 f t) = divMod f (Poly (G.fromList [-t, 1]))
178+
divModByDeg1 :: (Eq a, Num a, G.Vector vec a) => Poly vec a -> a -> (Poly vec a, a)
179+
divModByDeg1 f t = let w = G.postscanr (\a b -> a + b * t) 0 $ coeffAsc f
180+
in (Poly (G.tail w), G.head w)

0 commit comments

Comments
(0)

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