faster choose

1 files changed, 6 insertions, 4 deletions

diff --git a/bin/cluster.hs b/bin/cluster.hs
index 42e25dc..67b6e02 100644
--- a/bin/cluster.hs
+++ b/bin/cluster.hs
@@ -25,11 +25,13 @@ logFromInt = logFrom . fromIntegral
 logFrom = Exp . log
 
 binom :: Int -> Int -> Double -> Log Double
-binom 0 _ _ = 1
-binom n 0 p = logFrom (1 - p) ** logFromInt n
-binom n@(logFromInt -> n') k@(logFromInt -> k') p@(logFrom -> p') = choose * p' ** k' * logFrom (1 - p) ** (n' - k')
+binom n@(logFromInt -> n') k@(logFromInt -> k') p
+  | n == 0 = 1
+  | k == 0 = logFrom (1 - p) ** logFromInt n
+  | k == n = logFrom p ** logFromInt n
+  | otherwise = n `choose` k * logFrom p ** k' * logFrom (1 - p) ** (n' - k')
   where
-    choose = product [(n' + 1 - logFromInt i) / logFromInt i | i <- [1 .. k]]
+    choose (fromIntegral -> n) (fromIntegral -> k) = stirling n / stirling k / stirling (n - k)
 
 -- (infinite) binary trees
 data Tree a = Tree a (Tree a) (Tree a)