[issue37295] Possible optimizations for math.comb()

[Raymond Hettinger]

Raymond Hettinger <raymond.hettinger at gmail.com> added the comment:

It's a little late, but I had a thought that code could be made more general, slightly faster, and much easier to understand if the popcount logic were to be replaced with a table that records how many bits of the factorial were shifted out to make it odd.

from math import comb, perm, factorial as fact

Modulus = 2 ** 64
Cmax = 67
Pmax = 25
Fmax = Pmax

F = []          # odd factorial
S = []          # shift back to factorial
Finv = []       # multiplicative inverse of odd fact
for n in range(Cmax+1):
    f = fact(n)
    s = (f & -f).bit_length() - 1
    odd_f = (f >> s) % Modulus
    inv_f = pow(odd_f, -1, Modulus)
    assert odd_f * inv_f % Modulus == 1
    assert (odd_f << s) % Modulus == f % Modulus
    F.append(odd_f)
    S.append(s)
    Finv.append(inv_f)

def fact_small(n):
    return F[n] << S[n]

def perm_small(n, k):
    return (F[n] * Finv[n-k] % Modulus) << (S[n] - S[n-k])

def comb_small(n, k):
    return (F[n] * Finv[k] * Finv[n-k] % Modulus) << (S[n] - S[k] - S[n-k])

assert all(fact_small(n) == fact(n) for n in range(Fmax+1))
assert all(perm_small(n, k) == perm(n, k) for n in range(Pmax+1) for k in range(0, n+1))
assert all(comb_small(n, k) == comb(n, k) for n in range(Cmax+1) for k in range(0, n+1))

----------

_______________________________________
Python tracker <report at bugs.python.org>
<https://bugs.python.org/issue37295>
_______________________________________