Combinatorics problems are a major part of programming contests. This article explains how to efficiently count combinations using modular arithmetic.

Combinations

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

The formula for calculating all combinations without repetition is:

We will build an algorithm that efficiently precomputes factorials and evaluates this formula using modular arithmetic, which is a standard technique in contest problems to avoid handling very large numbers.

Factorials precomputation

Factorials grow extremely fast, so these problems typically ask for the answer modulo a prime number. Since we are using modular arithmetic, we can easily precompute all factorials up to n. That covers the numerator of the formula.

factorial = [1]

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

Multiplicative inverse

The challenge is computing the denominator. Modular arithmetic does not work with divisions, so we replace division with multiplication by 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. The theorem states that when p is a prime number, a^p - a is an integer multiple of p:

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

Dividing both sides by a, we get:

This gives us exactly the multiplicative inverse we need. In code:

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

We can then modify this 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 just as we did for factorials. We can compute all inverse factorials in descending order:

In code:

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

For simplicity, we put precomputation inside the final function comb(), which means it recomputes everything on each call. The precomputation should happen once outside the function.

Problems where the technique can be used:

LeetCode 3881 - Direction Assignments with Exactly K Visible People