Combinatorics problems are major part of programming contests. This article tries to explain how to efficiently count combinations with the help of modular arithmetic.

Combinations

A combination is a specific selection of k items from a set of size n. Combination can be either with or without repetition, but we will focus only on combinations without reperition which is frequently used in harder problems.

The formula for calculating all combinations without repetition is:

We are going to make an algorithm to efficiently precompute factorials and calculate the formula, using modular arithmetic though. That is a standard part of the problems to avoid big numbers

Factorials precomputation

Calculating factorials can very quickly lead to very big numbers, so these kinds of problems standardly ask for an answer modulo a prime number. Since we are using modular arithmetic, we can easily precompute all factorials up to n. That is enough for nominator part of the formula.

factorial = [1]

for i in range(1, n + 1):
    factorial.append(factorial[i - 1] * i % mod)

Multiplicative inverse

The tricky part is to calculate the denominator. Modular arithmetic does not work with divisions, so the first technique to use is the multiplicative inverse:

This way, we can express the denominator using inverse factorials.

Fermat's little theorem

To compute the multiplicative inverse, we use Fermat's little theorem. Theorem says that when p is a prime number, a^p - a is an integer multiple of p:

And if a is not divisible by p, then we get this statement:

Dividing both sides by a and rearranging, we can get this formula:

Combined with the modular arithmetic over primes discussed earlier, this is precisely a multiplicative inverse. Now we can put it into the code.

inverse_a = pow(a, p - 2, p)

And we can modify it to calculate inverse n! by using precomputed n! from the step before.

inverse_factorial[n] = pow(factorial[n], mod - 2, mod)

Inverse factorials precomputation

The next step is to precompute inverse factorials similarly to factorials above. With the help of maths, we can calculate them downwards:

inverse_factorial = [0] * (n + 1)
inverse_factorial[n] = pow(factorial[n], mod - 2, mod)

for i in range(n - 1, -1, -1):
    inverse_factorial[i] = (i + 1) * inverse_factorial[i + 1] % mod

Conclusion

Now that we have precomputed factorials and inverse factorials up to n, we can write the final function to calculate the total number of combinations without repetition for n and k.

def comb(n, k, mod):
    factorial = [1]

    for i in range(1, n + 1):
        factorial.append(factorial[i - 1] * i % mod)
    
    inverse_factorial = [0] * (n + 1)
    inverse_factorial[n] = pow(factorial[n], mod - 2, mod)
    
    for i in range(n - 1, -1, -1):
        inverse_factorial[i] = (i + 1) * inverse_factorial[i + 1] % mod

    return factorial[n] * inverse_factorial[k] % mod * inverse_factorial[n - k] % mod

Problems where the technique can be used:

LeetCode 3881 - Direction Assignments with Exactly K Visible People