{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
module Math.Combinatorics.Exact.Binomial (choose) where
import Data.List (foldl')
import Math.Combinatorics.Exact.Primes (primes)
choose :: (Integral a) => a -> a -> a
{-# SPECIALIZE choose ::
Integer -> Integer -> Integer,
Int -> Int -> Int
#-}
a
n choose :: a -> a -> a
`choose` a
k_
| a
n a -> Bool -> Bool
`seq` a
k_a -> Bool -> Bool
`seq` Bool
False = a
forall a. HasCallStack => a
undefined
| a
0 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
k_ Bool -> Bool -> Bool
&& a
k_ a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
n =
a
k a -> a -> a
`seq` a
nk a -> a -> a
`seq` a
sqrtN a -> a -> a
`seq`
(a -> Int -> a) -> a -> [Int] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
(\a
acc Int
prime -> a -> a -> a
step a
acc (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
prime))
a
1
((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>=) [Int]
primes)
| a
0 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
k_ Bool -> Bool -> Bool
&& a
k_ a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
n = a
1
| Bool
otherwise = a
0
where
k :: a
k = a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$! if a
k_ a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
n a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
2 then a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
k_ else a
k_
nk :: a
nk = a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
k
sqrtN :: a
sqrtN = Double -> a
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double
forall a. Floating a => a -> a
sqrt (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n) :: Double) a -> a -> a
forall a. a -> a -> a
`asTypeOf` a
n
step :: a -> a -> a
step a
acc a
prime
| a
acc a -> Bool -> Bool
`seq` a
prime a -> Bool -> Bool
`seq` Bool
False = a
forall a. HasCallStack => a
undefined
| a
prime a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
nk = a
acc a -> a -> a
forall a. Num a => a -> a -> a
* a
prime
| a
prime a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
n a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
2 = a
acc
| a
prime a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
sqrtN =
if a
n a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
prime a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
k a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
prime
then a
acc a -> a -> a
forall a. Num a => a -> a -> a
* a
prime
else a
acc
| Bool
otherwise = a
acc a -> a -> a
forall a. Num a => a -> a -> a
* a -> a -> a -> a -> a
go a
n a
k a
0 a
1
where
go :: a -> a -> a -> a -> a
go a
n' a
k' a
r a
p
| a
n' a -> Bool -> Bool
`seq` a
k' a -> Bool -> Bool
`seq` a
r a -> Bool -> Bool
`seq` a
p a -> Bool -> Bool
`seq` Bool
False = a
forall a. HasCallStack => a
undefined
| a
n' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 = a
p
| a
n' a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
prime a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< (a
k' a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
prime) a -> a -> a
forall a. Num a => a -> a -> a
+ a
r
= a -> a -> a -> a -> a
go (a
n' a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
prime) (a
k' a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
prime) a
1 (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$! a
p a -> a -> a
forall a. Num a => a -> a -> a
* a
prime
| Bool
otherwise = a -> a -> a -> a -> a
go (a
n' a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
prime) (a
k' a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
prime) a
0 a
p