counting combinations and permutations efficiently
I have some code to count permutations and combinations, and I'm trying to make it work better for large numbers.
I've found a better algorithm for permutations that avoids large intermediate results, but I still think I can do better for combinations.
So far, I've put in a special case to reflect the symmetry of nCr, but I'd still like to find a better algorithm that avoids the call to factorial(r), which is an unnecessarily large intermediate result. Without this optimization, the last doctest takes too long trying to calculate factorial(99000).
Can anyone suggest a more efficient way to count combinations?
from math import factorial
def product(iterable):
prod = 1
for n in iterable:
prod *= n
return prod
def npr(n, r):
"""
Calculate the number of ordered permutations of r items taken from a
population of size n.
>>> npr(3, 2)
6
>>> npr(100, 20)
1303995018204712451095685346159820800000
"""
assert 0 <= r <= n
return product(range(n - r + 1, n + 1))
def ncr(n, r):
"""
Calculate the number of unordered combinations of r items taken from a
population of size n.
>>> ncr(3, 2)
3
>>> ncr(100, 20)
535983370403809682970
>>> ncr(100000, 1000) == ncr(100000, 99000)
True
"""
assert 0 <= r <= n
if r > n // 2:
r = n - r
return npr(n, r) // factorial(r)
if n is not far from r then using the recursive definition of combination is probably better, since xC0 == 1 you will only have a few iterations:
The relevant recursive definition here is:
nCr = (n-1)C(r-1) * n/r
This can be nicely computed using tail recursion with the following list:
[(n - r, 0), (n - r + 1, 1), (n - r + 2, 2), ..., (n - 1, r - 1), (n, r)]
which is of course easily generated in Python (we omit the first entry since nC0 = 1) by izip(xrange(n - r + 1, n+1), xrange(1, r+1))
Note that this assumes r <= n you need to check for that and swap them if they are not. Also to optimize use if r < n/2 then r = n - r.
Now we simply need to apply the recursion step using tail recursion with reduce. We start with 1 since nC0 is 1 and then multiply the current value with the next entry from the list as below.
from itertools import izip
reduce(lambda x, y: x * y[0] / y[1], izip(xrange(n - r + 1, n+1), xrange(1, r+1)), 1)
Two fairly simple suggestions:
To avoid overflow, do everything in log space. Use the fact that log(a * b) = log(a) + log(b), and log(a / b) = log(a) - log(b). This makes it easy to work with very large factorials: log(n! / m!) = log(n!) - log(m!), etc.
Use the gamma function instead of factorial. You can find one in
scipy.stats.loggamma
. It's a much more efficient way to calculate log-factorials than direct summation.loggamma(n) == log(factorial(n - 1))
, and similarly,gamma(n) == factorial(n - 1)
.
There's a function for this in scipy which hasn't been mentioned yet: scipy.special.comb. It seems efficient based on some quick timing results for your doctest (~0.004 seconds for comb(100000, 1000, 1) == comb(100000, 99000, 1)
).
[While this specific question seems to be about algorithms the question is there a math ncr function in python is marked as a duplicate of this...]
If you don't need a pure-python solution, gmpy2 might help (gmpy2.comb
is very fast).
from scipy import misc
misc.comb(n, k)
should allow you to count combinations