Numbers of combinations modulo m, efficiently

Question:

First of all I’m solving a programming problem rather than a math problem now.

The question is

Anish got an unbiased coin and he tossed it n times and he asked Gourabh to count all the number of possible outcomes with j heads, for all j from 0 to n. Since the number of possible outcome can be huge, he will tell the values modulo m. To be clear, we need to return one integer per value of j.

The question is simple, but the problem arises with the time limit, being 1.5 seconds, but with input n as large as 200000.

I used math.comb to calculate the values, but it took more than 1.5 seconds to run.

So, are there any ways to calculate combinations in a faster way?

Edit#1:

Sample input:
2 998244353

Sample output: 1 2 1

Edit#2:

Here is the code that I’ve tried:

import math
 
n,m=input().split()
n = int(n)
m = int(m)
 
l = []
for i in range(n+1):
    l.append(math.comb(n,i)%m)
print(*l)

P.S: Please let me know if this is off topic for this site and suggest a suitable SE site to post this question. Thanks in advance! This question is from an inter college contest which ended 2 months ago.

Here is the original problem: https://codeforces.com/gym/430360/problem/B (you’ll need an account, and first time follow the "Contest Link" here to enter).

In case you are not able to view the problem, please find the picture below.

Asked By: Adithya

||

Answers:

I will sugest to use something like Modulo Multiplication.

If you have two numbers a and b, and you calculate their product modulo mod, you can take the modulo of each number first and then calculate the product. That is:

(a * b) % mod = ((a % mod) * (b % mod)) % mod

Taking into account that: C(n, r) = n! / (r! * (n - r)!)
You can probably get some nice reduction of computation passing mod inside a getting a general result.

Since you need to output values for all j, you should use this generating function:

If j >= 1, then C(n,j) = C(n,j-1) * (n+1-j) / j

Normally, when this sort of question is asked, the modulus m is a prime larger than n. This makes it relatively easy to do all these calculations mod m, because every j will have a multiplicative inverse.

In fact, it is so unusual to ask this question with a non-prime modulus, that I bet the codeforces problem description just fails to mention it. I would try it with the prime modulus assumption.

If you’re using python 3.8 or better, then there is a modular inverse built into the language, and you can do it just like this:

def getBinomialCoefficientsMod(n,m):
    result = [1]
    for j in range(1,n+1):
        result.append(( result[j-1] * (n+1-j) * pow(j,-1,m) )%m)
    return result

Edit: Well, it turns out that m is not always a large enough prime, and I don’t want to leave this answer incomplete, so here’s a version that works with composite or small m:

def getBinomialCoefficientsMod(n,m):

    # get the prime factors of the modulus
    facs=[]
    fac=2
    mleft = m
    while fac*fac<=mleft:
        if m%fac==0:
            facs.append(fac)
            while mleft%fac==0:
                mleft//=fac
        fac+=1
    if mleft>1:
        facs.append(mleft)

    result = [1]
    # factor of the last result that is relatively prime to m
    rpresult = 1
    # powers of the prime factors of m in the last result
    exponents = [0]*len(facs)

    for j in range(1,n+1):
        p=1
        num = n+1-j
        den = j

        # remove factors of the modulus from num and den,
        # track their exponents, and get their product
        for i in range(len(facs)):
            fac = facs[i]
            while num%fac==0:
                exponents[i]+=1
                num//=fac
            while den%fac==0:
                exponents[i]-=1
                den//=fac
            p = p*pow(fac,exponents[i],m)
    
        rpresult = (rpresult * num * pow(den,-1,m)) % m
    
        result.append(( rpresult * p )%m)

    return result

Answered By: Matt Timmermans

Using the usual multiplicative formula to compute the next number from the previous, but with keeping the numbers small. Let’s first look at a naive version for clarity.

Naive

def naive(n, m):
    c = 1
    yield c
    for k in range(n):
        c = c * (n-k) // (k+1)
        yield c % m

n, m = map(int, input().split())
print(*naive(n, m))

Takes me ~30 seconds with n=200000. Because c grows very large, up to 60204 digits (199991 bits). And calculations with such large numbers are slow.

Fast

Instead of naively computing those large c and using modulo m only for output, let’s keep c small throughout, modulo m. Got accepted on the site, taking ~0.68 seconds.

from math import gcd

def fast(n, m):
    c = 1
    G = 1
    yield c
    for k in range(n):

        mul = n - k
        while (g := gcd(mul, m)) > 1:
            mul //= g
            G *= g

        div = k + 1
        while (g := gcd(div, m)) > 1:
            div //= g
            G //= g

        c = c * mul * pow(div, -1, m) % m
        yield c * G % m

n, m = map(int, input().split())
print(*fast(n, m))

Attempt This Online!

Multiplication is fine under modulo. If it were only c = c * (n-k), we could just do c = c * (n-k) % m.

Division doesn’t allow that. So instead of dividing by k+1, we multiply with its inverse (k+1)-1 modulo m. The inverse of some number x is the number x-1 so you get x·x-1 = 1. For example, 7-1 modulo 10 is 3. Because multiplying 7 and 3 gives you 21, which is 1 (modulo 10).

Next issue: Not all numbers have an inverse modulo m. For example, 6 doesn’t have an inverse modulo 10. You can’t multiply 6 with any integer and get 1 (modulo 10). Because 6 and 10 have common divisor 2. What we’ll do is invert as much of 6 as possible. Extract the common divisor 2, leaving us with 3. That does have an inverse modulo 10 (namely 7).

So extract prime factors in the multipliers/divisors common with m into a separate number G. And update c with what remains, modulo m. Then combine c and G for output.

Rough times I get for n=200000, m=998244353 (the large prime from the question):

naive: 30.0 seconds
fast:   1.0 seconds
Matt's: 1.0 seconds

For n=200000, m=2*3*5*7*11*13*17*19*23:

naive: 30.0 seconds
fast:   1.2 seconds
Matt's: 4.8 seconds

I think worst case is a modulus with many primes like m=2*3*5*7*11*13*17*19*23, that maximizes my G. With n=200000, G grows up to 127 bits. Nothing to worry about.

My solution/explanation for a similar problem on Leetcode. That had modulus 10 and I hardcoded factors 2 and 5 and counted them instead of multiplying them into a number G like I did here. Maybe I’ll revisit it with this general solution…

Answered By: Stefan Pochmann