fork download
  1. {-# LANGUAGE Rank2Types #-}
  2. {-# LANGUAGE BangPatterns #-}
  3.  
  4. import Data.List
  5. import Control.Applicative
  6. import Control.Monad
  7. import Control.Arrow
  8.  
  9. _u = id
  10.  
  11. dup :: a -> (a, a)
  12. dup x = (x, x)
  13.  
  14. infixr 9 .*
  15. (.*) :: (c -> d) -> (a -> b -> c) -> a -> b -> d
  16. (.*) = (.) . (.)
  17.  
  18. (!-) :: Integer -> Integer -> Integer
  19. n !- m = let d = n - m in if d <= 0 then 0 else d
  20.  
  21. addToAcc :: c' -> (c, b) -> ((c, c'), b)
  22. addToAcc x (acc, res) = ((acc, x), res)
  23.  
  24. infixr 5 $: , $++
  25. infixr 9 $!!
  26.  
  27. newtype AccList a = AccList {runAccList :: forall b c.
  28. Integer -> c -> b -> (c -> a -> (c, b -> b)) -> b
  29. }
  30.  
  31. instance Functor AccList where
  32. fmap f xs = pure f <*> xs
  33.  
  34. instance Applicative AccList where
  35. pure = return
  36. (<*>) = ap
  37.  
  38. instance Monad AccList where
  39. return = ($: nil)
  40. (>>=) = flip fconcatMap
  41.  
  42. nil :: AccList a
  43. nil = AccList $ \_ _ res _ -> res
  44.  
  45. ($:) :: a -> AccList a -> AccList a
  46. x $: xs = AccList $ \d acc res facc -> if d == 0
  47. then let (acc', fres) = facc acc x in
  48. fres $ runAccList xs 0 acc' res facc
  49. else runAccList xs (d - 1) acc res facc
  50.  
  51. runAccListAccEnd :: AccList a -> forall b c.
  52. Integer -> c -> (c -> (c, b)) -> (c -> a -> (c, b -> b)) -> (c, b)
  53. runAccListAccEnd xs d acc end facc = ($ acc) $ runAccList xs d acc end $
  54. \(!acc') x -> let (acc'', fres) = facc acc' x in (,) acc'' $
  55. \end' -> let (acc''', res') = end' acc'' in
  56. const (acc''', fres res')
  57.  
  58. runAccListAcc :: AccList a -> forall b c.
  59. Integer -> c -> b -> (c -> a -> (c, b -> b)) -> (c, b)
  60. runAccListAcc xs d acc res facc =
  61. runAccListAccEnd xs d acc (flip (,) res) facc
  62.  
  63. runAccListEnd :: AccList a -> forall b c.
  64. Integer -> c -> (c -> b) -> (c -> a -> (c, b -> b)) -> b
  65. runAccListEnd xs d acc end facc =
  66. snd $ runAccListAccEnd xs d acc (second end . dup) facc
  67.  
  68. ffromList = foldr ($:) nil
  69. ftoList = ffoldr (:) []
  70.  
  71. -- Basic functions.
  72.  
  73. ($++) :: AccList a -> AccList a -> AccList a
  74. xs $++ ys = AccList $ \d acc res facc ->
  75. let (acc', res') = runAccListAcc xs d acc
  76. (runAccList ys (d !- flength xs) acc' res facc) facc
  77. in res'
  78.  
  79. fhead :: AccList a -> a
  80. fhead = ffoldr const (error "empty list")
  81.  
  82. flast :: AccList a -> a
  83. flast xs = fst $ runAccListAcc xs 0 (fhead xs) () $ \_ x -> (x, _u)
  84.  
  85. ftail :: AccList a -> AccList a
  86. ftail = fdrop 1
  87.  
  88. finit :: AccList a -> AccList a
  89. finit xs = ftake (flength xs - 1) xs
  90.  
  91. fnull :: AccList a -> Bool
  92. fnull = ffoldr (\_ _ -> True) False
  93.  
  94. flength :: AccList a -> Integer
  95. flength = ffoldl' (\r _ -> r + 1) 0
  96.  
  97. -- List transformations
  98.  
  99. ffmap :: (a -> b) -> AccList a -> AccList b
  100. ffmap f xs = AccList $ \d acc res facc ->
  101. runAccList xs d acc res $ \acc' -> facc acc' . f
  102.  
  103. freverse :: AccList a -> AccList a
  104. freverse = ffoldl (flip ($:)) nil
  105.  
  106. -- fintersperse
  107.  
  108. -- fintercalate
  109.  
  110. -- ftranspose
  111.  
  112. fsubsequences :: AccList a -> AccList (AccList a)
  113. fsubsequences = (nil $:) . f (\x -> (pure x $:) . f (\ys r -> ys $: (x $: ys) $: r))
  114. where f = flip ffoldr nil
  115.  
  116. -- fpermutations
  117.  
  118. -- Reducing lists (folds)
  119.  
  120. ffoldl :: (b -> a -> b) -> b -> AccList a -> b
  121. ffoldl f z xs = fst $ runAccListAcc xs 0 z () $ \res' x -> (f res' x, id)
  122.  
  123. ffoldl' :: (b -> a -> b) -> b -> AccList a -> b
  124. ffoldl' f z xs = fst $ runAccListAcc xs 0 z () $ \(!res') x -> (f res' x, id)
  125.  
  126. -- ffoldl1
  127.  
  128. -- ffoldl1'
  129.  
  130. ffoldr :: (a -> b -> b) -> b -> AccList a -> b
  131. ffoldr f z xs = runAccList xs 0 () z (\() x -> ((), f x))
  132.  
  133. -- Special folds
  134.  
  135. fconcat :: AccList (AccList a) -> AccList a
  136. fconcat = ffoldr ($++) nil
  137.  
  138. fconcatMap :: (a -> AccList b) -> AccList a -> AccList b
  139. fconcatMap = fconcat .* ffmap
  140.  
  141. fand :: AccList Bool -> Bool
  142. fand = fall id
  143.  
  144. for :: AccList Bool -> Bool
  145. for = fany id
  146.  
  147. fany :: (a -> Bool) -> AccList a -> Bool
  148. fany p = ffoldr ((||) . p) False
  149.  
  150. fall :: (a -> Bool) -> AccList a -> Bool
  151. fall p = ffoldr ((&&) . p) True
  152.  
  153. fsum :: Num a => AccList a -> a
  154. fsum = ffoldl' (+) 0
  155.  
  156. fproduct :: Num a => AccList a -> a
  157. fproduct = ffoldl' (*) 1
  158.  
  159. -- fmaximum
  160.  
  161. -- fminimum
  162.  
  163. -- Building lists
  164.  
  165. -- Scans
  166.  
  167. fscanl :: (b -> a -> b) -> b -> AccList a -> AccList b
  168. fscanl f z xs = AccList $ \d acc res facc ->
  169. let (z', _) = runAccListAcc (ftake d xs) 0 z res
  170. (\z'' x -> (f z'' x, _u))
  171. in runAccListEnd xs d (acc, z')
  172. (\(acc', z'') -> ($ res) $ snd $ facc acc' z'')
  173. (\(acc', z'') x -> addToAcc (f z'' x) $ facc acc' z'')
  174.  
  175. fscanl1 :: (a -> a -> a) -> AccList a -> AccList a
  176. fscanl1 f xs = fscanl f (fhead xs) (ftail xs)
  177.  
  178. fscanr :: (a -> b -> b) -> b -> AccList a -> AccList b
  179. fscanr f = uncurry ($:) .* fmapAccumR (dup .* flip f)
  180.  
  181. -- fscanr1
  182.  
  183. -- Accumulating maps
  184.  
  185. fmapAccumL :: (c -> a -> (c, b)) -> c -> AccList a -> (c, AccList b)
  186. fmapAccumL f z xs = runAccListAcc xs 0 z nil $ second ($:) .* f
  187.  
  188. fsndMapAccumL :: (c -> a -> (c, b)) -> c -> AccList a -> AccList b
  189. fsndMapAccumL f z xs = AccList $ \d acc res facc ->
  190. let (z', _) = runAccListAcc (ftake d xs) 0 z res $
  191. \z'' x -> (fst $ f z'' x, _u)
  192. in runAccList xs d (acc, z') res $
  193. \(acc', z'') x -> let (z''', x') = f z'' x in
  194. addToAcc z''' $ facc acc' x'
  195.  
  196. fmapAccumR :: (c -> a -> (c, b)) -> c -> AccList a -> (c, AccList b)
  197. fmapAccumR f z = ffoldr (\x (z', xs') -> ($: xs') `second` f z' x) (z, nil)
  198.  
  199. -- Infinite lists
  200.  
  201. frepeat :: a -> AccList a
  202. frepeat = ffromList . repeat
  203.  
  204. fiterate :: (a -> a) -> a -> AccList a
  205. fiterate = ffromList .* iterate
  206.  
  207. freplicate :: Integer -> a -> AccList a
  208. freplicate n = ftake n . frepeat
  209.  
  210. fcycle :: AccList a -> AccList a
  211. fcycle xs | fnull xs = error "Empty list"
  212. fcycle xs = xs' where xs' = xs $++ xs'
  213.  
  214. -- Unfolding
  215.  
  216. funfoldr :: (b -> Maybe (a, b)) -> b -> AccList a
  217. funfoldr = ffromList .* unfoldr
  218.  
  219. -- Sublists
  220.  
  221. -- Extracting sublists
  222.  
  223. ftake :: Integer -> AccList a -> AccList a
  224. ftake n xs = AccList $ \d acc res facc ->
  225. runAccList xs d (acc, n - d) res $
  226. \(acc', n') x -> if n' <= 0
  227. then ((acc', 0), const res)
  228. else addToAcc (n' - 1) $ facc acc' x
  229.  
  230. fdrop :: Integer -> AccList a -> AccList a
  231. fdrop n xs = AccList $ \d -> runAccList xs (d + (n !- 0))
  232.  
  233. fsplitAt :: Integer -> AccList a -> (AccList a, AccList a)
  234. fsplitAt n xs = (ftake n xs, fdrop n xs)
  235.  
  236. ftakeWhile :: (a -> Bool) -> AccList a -> AccList a
  237. ftakeWhile p xs = AccList $ \d acc res facc ->
  238. if fall p (ftake d xs)
  239. then runAccList xs d acc res $
  240. \acc' x -> if p x then facc acc' x else (acc', const res)
  241. else res
  242.  
  243. fdropWhile :: (a -> Bool) -> AccList a -> AccList a
  244. fdropWhile p xs = fdrop (flength (ftakeWhile p xs)) xs
  245.  
  246. -- fdropWhileEnd
  247.  
  248. fspan :: (a -> Bool) -> AccList a -> (AccList a, AccList a)
  249. fspan p xs = (ftakeWhile p xs, fdropWhile p xs)
  250.  
  251. fbreak :: (a -> Bool) -> AccList a -> (AccList a, AccList a)
  252. fbreak p = fspan (not . p)
  253.  
  254. -- fstripPrefix
  255.  
  256. -- fgroup
  257.  
  258. finits :: AccList a -> AccList (AccList a)
  259. finits = ffoldr (\x r -> nil $: fmap (x $:) r) (nil $: nil)
  260.  
  261. ftails :: AccList a -> AccList (AccList a)
  262. ftails = ffoldr (\x r -> (x $: fhead r) $: r) (nil $: nil)
  263.  
  264. -- Predicates
  265.  
  266. -- fisPrefixOf
  267.  
  268. -- fisSuffixOf
  269.  
  270. -- fisInfixOf
  271.  
  272. -- Searching lists
  273.  
  274. -- Searching by equality
  275.  
  276. felem :: Eq a => a -> AccList a -> Bool
  277. felem = fany . (==)
  278.  
  279. fnotElem :: Eq a => a -> AccList a -> Bool
  280. fnotElem = fall . (/=)
  281.  
  282. flookup :: Eq a => a -> AccList (a, b) -> Maybe b
  283. flookup k = ffoldr (\(k', x) r -> if k == k' then Just x else r) Nothing
  284.  
  285. -- Searching with a predicate
  286.  
  287. ffind :: (a -> Bool) -> AccList a -> Maybe a
  288. ffind = flistToMaybe .* ffilter
  289.  
  290. ffilter :: (a -> Bool) -> AccList a -> AccList a
  291. ffilter p xs = AccList $ \d acc res facc -> runAccList xs d acc res $
  292. \acc' x -> if p x then facc acc' x else (acc', id)
  293.  
  294. fpartition :: (a -> Bool) -> AccList a -> (AccList a, AccList a)
  295. fpartition p xs = (ffilter p xs, ffilter (not . p) xs)
  296.  
  297. -- fIndexing lists
  298.  
  299. ($!!) :: AccList a -> Integer -> a
  300. xs $!! n = fhead $ fdrop n xs
  301.  
  302. felemIndex :: Eq a => a -> AccList a -> Maybe Integer
  303. felemIndex = ffindIndex . (==)
  304.  
  305. felemIndices :: Eq a => a -> AccList a -> AccList Integer
  306. felemIndices = ffindIndices . (==)
  307.  
  308. ffindIndex :: (a -> Bool) -> AccList a -> Maybe Integer
  309. ffindIndex = flistToMaybe .* ffindIndices
  310.  
  311. ffindIndices :: (a -> Bool) -> AccList a -> AccList Integer
  312. ffindIndices p = fmap fst . ffilter (p . snd) . fenumerate
  313.  
  314. -- Maybe
  315.  
  316. flistToMaybe :: AccList a -> Maybe a
  317. flistToMaybe = ffoldr (\x _ -> Just x) Nothing
  318.  
  319. -- Misc
  320.  
  321. fenumerate :: AccList a -> AccList (Integer, a)
  322. fenumerate = fsndMapAccumL (\n x -> (n + 1, (n, x))) 0
  323.  
  324. ---
  325.  
  326. fdefhead = ffoldr const
  327.  
  328. fpow n f x = iterate f x !! n
  329.  
  330. -- quadratic
  331. fzipWith f xs ys = fsndMapAccumL
  332. (\x's' y' -> (ftail x's', alt f (fhead x's') y')) x's y's
  333. where (alt, x's, y's) = if flength xs < flength ys
  334. then (flip, ys, xs)
  335. else (id, xs, ys)
  336.  
  337. main = do
  338. print $ sum $ fpow 5000 tail [1..10^6]
  339. print $ fsum $ fpow 5000 ftail $ ffromList [1..10^6]
  340.  
  341. print $ take 10 $ drop 5 $ take 20 $ take 27 $ findIndices even $ drop 5 $ scanl (-) 2 $ take 40 $ drop 10 [1..]
  342. print $ ftoList $ ftake 10 $ fdrop 5 $ ftake 20 $ ftake 27 $ ffindIndices even $ fdrop 5 $ fscanl (-) 2 $ ftake 40 $ fdrop 10 $ ffromList [1..]
  343.  
  344. print $ reverse $ take 10 $ drop 3 $ reverse $ take 28 $ filter even $ drop 5 $ snd $ mapAccumL (\x y -> (x+y, x*y)) 5 $ takeWhile (<= 30) $ take 40 $ drop 10 [1..]
  345. print $ ftoList $ freverse $ ftake 10 $ fdrop 3 $ freverse $ ftake 28 $ ffilter even $ fdrop 5 $ fsndMapAccumL (\x y -> (x+y, x*y)) 5 $ ftakeWhile (<= 30) $ ftake 40 $ fdrop 10 $ ffromList [1..]
  346.  
  347. print $ sum $ reverse $ drop 60 $ reverse $ takeWhile (<= 10^3) $ take (10^10) $ dropWhile (<= 50) $ drop 30 [1..]
  348. print $ fsum $ freverse $ fdrop 60 $ freverse $ ftakeWhile (<= 10^3) $ ftake (10^10) $ fdropWhile (<= 50) $ fdrop 30 $ ffromList [1..]
  349.  
  350. print $ drop 2 $ takeWhile (<= 6) $ concat $ subsequences $ dropWhile (<= 3) [1..10]
  351. print $ ftoList $ fdrop 2 $ ftakeWhile (<= 6) $ fconcat $ fsubsequences $ fdropWhile (<= 3) $ ffromList [1..10]
Success #stdin #stdout 0.49s 7288KB
stdin
Standard input is empty
stdout
499987997500
499987997500
[11,14,15,18,19,22,23,26,27,30]
[11,14,15,18,19,22,23,26,27,30]
[1120,1462,1854,2800,3360,3982,5424,6250,7150]
[1120,1462,1854,2800,3360,3982,5424,6250,7150]
440995
440995
[4,5,6,4,6,5,6,4,5,6]
[4,5,6,4,6,5,6,4,5,6]