{-# LANGUAGE TypeApplications #-}
module Math.NumberTheory.Recurrences.Pentagonal
( partition
) where
import qualified Data.Chimera as Ch
import Data.Vector (Vector)
import Numeric.Natural (Natural)
pents :: (Enum a, Num a) => [a]
pents :: [a]
pents = [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave ((a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (\a
acc a
n -> a
acc a -> a -> a
forall a. Num a => a -> a -> a
+ a
3 a -> a -> a
forall a. Num a => a -> a -> a
* a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a
0 [a
1..])
((a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (\a
acc a
n -> a
acc a -> a -> a
forall a. Num a => a -> a -> a
+ a
3 a -> a -> a
forall a. Num a => a -> a -> a
* a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
2) a
1 [a
2..])
where
interleave :: [a] -> [a] -> [a]
interleave :: [a] -> [a] -> [a]
interleave (a
n : [a]
ns) (a
m : [a]
ms) = a
n a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a
m a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave [a]
ns [a]
ms
interleave [a]
_ [a]
_ = []
partitionF :: Num a => (Word -> a) -> Word -> a
partitionF :: (Word -> a) -> Word -> a
partitionF Word -> a
_ Word
0 = a
1
partitionF Word -> a
f Word
n
= [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
([a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Num a => a -> a -> a
(*) ([a] -> [a]
forall a. [a] -> [a]
cycle [a
1, a
1, -a
1, -a
1])
([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ (Word -> a) -> [Word] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Word -> a
f (Word -> a) -> (Word -> Word) -> Word -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
-))
([Word] -> [a]) -> [Word] -> [a]
forall a b. (a -> b) -> a -> b
$ (Word -> Bool) -> [Word] -> [Word]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
n)
([Word] -> [Word]) -> [Word] -> [Word]
forall a b. (a -> b) -> a -> b
$ [Word] -> [Word]
forall a. [a] -> [a]
tail [Word]
forall a. (Enum a, Num a) => [a]
pents
partition :: Num a => [a]
partition :: [a]
partition = Chimera Vector a -> [a]
forall (v :: * -> *) a. Vector v a => Chimera v a -> [a]
Ch.toList (Chimera Vector a -> [a]) -> Chimera Vector a -> [a]
forall a b. (a -> b) -> a -> b
$ ((Word -> a) -> Word -> a) -> Chimera Vector a
forall (v :: * -> *) a.
Vector v a =>
((Word -> a) -> Word -> a) -> Chimera v a
Ch.tabulateFix @Vector (Word -> a) -> Word -> a
forall a. Num a => (Word -> a) -> Word -> a
partitionF
{-# SPECIALIZE partition :: [Int] #-}
{-# SPECIALIZE partition :: [Word] #-}
{-# SPECIALIZE partition :: [Integer] #-}
{-# SPECIALIZE partition :: [Natural] #-}