Possible optimizations for math.comb()

[rhettinger]

BPO	37295
Nosy	@tim-one, @rhettinger, @mdickinson, @serhiy-storchaka, @PedanticHacker, @mcognetta, @pochmann
PRs	bpo-37295: Add fast path for small n in math.comb #29020 bpo-37295: Optimize math.comb() for small arguments #29030 bpo-37295: Optimize math.comb() and math.perm() #29090 bpo-37295: Speed up math.comb(n, k) for 0 <= k <= n <= 67 #30275 bpo-37295: Use constant-time comb() for larger n depending on k #30305 bpo-37295: More direct computation of power-of-two factor in math.comb #30313
Files	comb_with_primes.py comb_with_side_limits.py: Compute limits for comb_small timecomb.py driver.py comb_pole.py: Experiment with precomputed diagonal -- runs on 3.8 or later comb_pole2.py: Precomputed k diagonal and handling for small k

^{Note: these values reflect the state of the issue at the time it was migrated and might not reflect the current state.}

Show more details

GitHub fields:

assignee = None
closed_at = None
created_at = <Date 2019-06-15.19:17:40.265>
labels = ['library', 'performance']
title = 'Possible optimizations for math.comb()'
updated_at = <Date 2022-01-23.22:42:15.922>
user = 'https://github.com/rhettinger'

bugs.python.org fields:

activity = <Date 2022-01-23.22:42:15.922>
actor = 'tim.peters'
assignee = 'none'
closed = False
closed_date = None
closer = None
components = ['Library (Lib)']
creation = <Date 2019-06-15.19:17:40.265>
creator = 'rhettinger'
dependencies = []
files = ['50439', '50526', '50530', '50531', '50557', '50559']
hgrepos = []
issue_num = 37295
keywords = ['patch']
message_count = 73.0
messages = ['345711', '345731', '345761', '345784', '345876', '345952', '345956', '404184', '404208', '404429', '404472', '406300', '406337', '406338', '406341', '407736', '408867', '408952', '408954', '408962', '408970', '408984', '408989', '408994', '409000', '409001', '409002', '409010', '409024', '409028', '409062', '409095', '409110', '409116', '409154', '409254', '409268', '409269', '409270', '409274', '409277', '409298', '409315', '409317', '409320', '409321', '409331', '409346', '409360', '409373', '409376', '409377', '409382', '409383', '409393', '409417', '409432', '410145', '410339', '410347', '410382', '410383', '410387', '410420', '410448', '410460', '410543', '410587', '410925', '410931', '411343', '411427', '411430']
nosy_count = 7.0
nosy_names = ['tim.peters', 'rhettinger', 'mark.dickinson', 'serhiy.storchaka', 'PedanticHacker', 'mcognetta', 'Stefan Pochmann']
pr_nums = ['29020', '29030', '29090', '30275', '30305', '30313']
priority = 'normal'
resolution = None
stage = 'patch review'
status = 'open'
superseder = None
type = 'performance'
url = 'https://bugs.python.org/issue37295'
versions = []

[rhettinger]

The implementation of math.comb() is nice in that it limits the size of intermediate values as much as possible and it avoids unnecessary steps when min(k,n-k) is small with respect to k.

There are some ways it can be made better or faster:

1. For small values of n, there is no need to use PyLong objects at every step. Plain C integers would suffice and we would no longer need to make repeated allocations for all the intermediate values. For 32-bit builds, this could be done for n<=30. For 64-bit builds, it applies to n<=62. Adding this fast path would likely give a 10x to 50x speed-up.

2. For slightly larger values of n, the computation could be sped-up by precomputing one or more rows of binomial coefficients (perhaps for n=75, n=150, and n=225). The calculation could then start from that row instead of from higher rows on Pascal's triangle.

For example comb(100, 55) is currently computed as:
comb(55,1) * 56 // 2 * 57 // 3 ... 100 // 45 <== 45 steps

Instead, it could be computed as:
comb(75,20) * 76 // 21 * 77 // 22 ... 100 / 45 <== 25 steps
^-- found by table lookup

This gives a nice speed-up in exchange for a little memory in an array of constants (for n=75, we would need an array of length 75//2 after exploiting symmetry). Almost all cases would should show some benefit and in favorable cases like comb(76, 20) the speed-up would be nearly 75x.

3. When k is close to n/2, the current algorithm is slower than just computing (n!) / (k! * (n-k)!). However, the factorial method comes at the cost of more memory usage for large n. The factorial method consumes memory proportional to n*log2(n) while the current early-cancellation method uses memory proportional to n+log2(n). Above some threshold for memory pain, the current method should always be preferred. I'm not sure the factorial method should be used at all, but it is embarrassing that factorial calls sometimes beat the current C implementation:

    $ python3.8 -m timeit -r 11 -s 'from math import comb, factorial as fact' -s 'n=100_000' -s 'k = n//2' 'comb(n, k)'
    1 loop, best of 11: 1.52 sec per loop
    $ python3.8 -m timeit -r 11 -s 'from math import comb, factorial as fact' -s 'n=100_000' -s 'k = n//2' 'fact(n) // (fact(k) * fact(n-k))'
    1 loop, best of 11: 518 msec per loop

4. For values such as n=1_000_000 and k=500_000, the running time is very long and the routine doesn't respond to SIGINT. We could add checks for keyboard interrupts for large n. Also consider releasing the GIL.

5. The inner-loop current does a pure python subtraction than could in many cases be done with plain C integers. When n is smaller than maxsize, we could have a code path that replaces "PyNumber_Subtract(factor, _PyLong_One)" with something like "PyLong_FromUnsignedLongLong((unsigned long long)n - i)".

[rhettinger]

Optimizations 2 and 3 look something like this:

    bc75 = [comb(75, r) for r in range(75//2+1)]
    bc150 = [comb(150, r) for r in range(150//2+1)]
    bc225 = [comb(225, r) for r in range(225//2+1)]

    def comb(n, k):
        if n < 0 or k < 0: raise ValueError
        if k > n: return 0
        k = min(k, n-k)
        if k > n // 3 and n < 100_000:
            return factorial(n) // (factorial(r) * factorial(n-r))
        if 75 <= n <= 75 + k:
            c, num, den, times = bc75[75-(n-k)], 75+1, 75-(n-k)+1, n-75
        elif 150 <= n <= 150 + k:
            c, num, den, times = bc150[150-(n-k)], 150+1, 150-(n-k)+1, n-150
        elif 225 <= n <= 225 + k:
            c, num, den, times = bc225[225-(n-k)], 225+1, 225-(n-k)+1, n-225
        else:
            c, num, den, times = 1, n-k+1, 1, k
        for i in range(times):
            c = c * num // den
            num += 1
            den += 1
        return c

[mdickinson]

(1), (4) and (5) sound good to me.

For (1), it might make sense to ignore the 32-bit vs. 64-bit distinction and use uint64_t for the internal computations. Then we can do up to n = 62 regardless of platform.

(2) feels like too much extra complication to me, but that would become clearer with an implementation.

For (3), I somewhat agree that the factorial method should be avoided.

For (4), I don't see how/when the GIL could be released: doesn't the algorithm involve lots of memory allocations/deallocations and reference count adjustments?

Can the suggested performance improvements go into 3.8, or should they wait for 3.9? It's not clear to me whether a performance improvement after feature freeze is okay or not.

[PedanticHacker]

Performance improvements is what a beta build exists for in the first place.

[rhettinger]

FWIW, here's a rough implementation of (3). It is limited to a specific range to avoid excess memory usage. For comb(50_000, 25_000), it gives a three-fold speedup:

    if (k_long > 20 && k_long > n_long / 3 && n_long <= 100000) {
        PyObject *den = math_factorial(module, k);                 /* den = k! */
        temp = PyNumber_Subtract(n, k);
        Py_SETREF(temp, math_factorial(module, temp));
        Py_SETREF(den, PyNumber_Multiply(den, temp));              /* den *= (n - k)! */
        Py_DECREF(temp);
        Py_SETREF(result, math_factorial(module, n));              /* result = n! */
        Py_SETREF(result, PyNumber_FloorDivide(result, den));      /* result //= (n-k)! */
        Py_DECREF(den);
        return result;
    }

Can the suggested performance improvements go into 3.8, or should they wait for 3.9?
It's not clear to me whether a performance improvement after feature freeze is okay or not.

Historically, we've used the beta phase for optimizations, tweaking APIs, and improving docs. However, I'm in no rush and this can easily wait for 3.9.

My only concern is that the other math functions, except for factorial(), have bounded running times and memory usage, so performance is more of a concern for this function which could end-up being an unexpected bottleneck or being a vector for a DOS attack. That said, we haven't had any negative reports regarding factorial(), so this may be a low priority.

[tim-one]

In real life, I expect 99.999%+ of calls will be made with small arguments, so (1) is worth it. I like Mark's suggestion to use uint64_t so the acceptable range doesn't depend on platform. At least in the world I live in, 32-bit boxes are all but extinct anyway.

I honestly wouldn't bother with more than that. It's fun to optimize giant-integer algorithms with an ever-ballooning number of different clever approaches, but Python is an odd place for that. People looking for blazing fast giant-int facilities generally want lots & lots of them, so are better steered toward, e.g., gmp. That's its reason for existing.

For example, their implementation of binomial coefficients uses special division algorithms exploiting that the quotient is exact (no remainder):

https://gmplib.org/manual/Exact-Division.html#Exact-Division

There's just no end to potential speedups. But in Python, I expect a vast majority of users will be happy if they get the right answer for the number of possible poker hands ;-)

Oh ya - some smart kid will file a bug report about the inability to interrupt the calculation of a billion-bit result, so (4) is inevitable. Me, I'd wait for them to complain, and encourage _them_ to learn something useful by writing a patch to fix it :-)

[serhiy-storchaka]

Me, I'd wait for them to complain, and encourage _them_ to learn something useful by writing a patch to fix it :-)

+1!

[serhiy-storchaka]

Here is more optimized PR inspired by PR 29020. It would be too long to explain how PR 29020 can be improved, so I write a new PR.

Basically it implements Raymond's idea #1, but supports n>62 for smaller k.

How to calculate limits:

import math
n = m = 2**64
k = 1
while True:
    nmax = int(math.ceil((m * math.factorial(k-1)) ** (1/k) + (k-1)/2)) + 100
    n = min(n, nmax)
    while math.comb(n, k) * k >= m:
        n -= 1
    if n < 2*k: break
    print(k, n)
    k += 1