module AOC.Common.Numeric (
    PascalTable(..)
  , buildPascalTable
  , binom
  , pascals
  ) where

import qualified Data.Vector.Unboxed as VU
import           Data.List (scanl')

pascals :: [[Int]]
pascals :: [[Int]]
pascals = Int -> [Int]
forall a. a -> [a]
repeat Int
1 [Int] -> [[Int]] -> [[Int]]
forall a. a -> [a] -> [a]
: ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map ([Int] -> [Int]
forall a. [a] -> [a]
tail ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0) [[Int]]
pascals

data PascalTable = PascalTable 
    { PascalTable -> Int
ptWidth :: !Int
    , PascalTable -> Int
ptDepth :: !Int
    , PascalTable -> Vector Int
ptTable :: !(VU.Vector Int)
    }
  deriving Int -> PascalTable -> ShowS
[PascalTable] -> ShowS
PascalTable -> String
(Int -> PascalTable -> ShowS)
-> (PascalTable -> String)
-> ([PascalTable] -> ShowS)
-> Show PascalTable
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PascalTable] -> ShowS
$cshowList :: [PascalTable] -> ShowS
show :: PascalTable -> String
$cshow :: PascalTable -> String
showsPrec :: Int -> PascalTable -> ShowS
$cshowsPrec :: Int -> PascalTable -> ShowS
Show

buildPascalTable
    :: Int            -- ^ num diagonals
    -> Int            -- ^ depth of diagonal
    -> PascalTable
buildPascalTable :: Int -> Int -> PascalTable
buildPascalTable Int
n Int
d = PascalTable :: Int -> Int -> Vector Int -> PascalTable
PascalTable
    { ptWidth :: Int
ptWidth = Int
n
    , ptDepth :: Int
ptDepth = Int
d
    , ptTable :: Vector Int
ptTable = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList ([Int] -> Vector Int)
-> ([[Int]] -> [Int]) -> [[Int]] -> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Int]] -> Vector Int) -> [[Int]] -> Vector Int
forall a b. (a -> b) -> a -> b
$
        Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
d ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [[Int]] -> [[Int]]
forall a. Int -> [a] -> [a]
take Int
n [[Int]]
pascals
    }

binom
    :: PascalTable
    -> Int          -- ^ diagonal
    -> Int          -- ^ index
    -> Int
binom :: PascalTable -> Int -> Int -> Int
binom PascalTable{Int
Vector Int
ptTable :: Vector Int
ptDepth :: Int
ptWidth :: Int
ptTable :: PascalTable -> Vector Int
ptDepth :: PascalTable -> Int
ptWidth :: PascalTable -> Int
..} Int
i Int
j = Vector Int
ptTable Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
ptDepth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)