Hi everyone!

There is an article on cp-algorithms about how to compute $$$n!$$$ modulo prime number $$$p$$$. Today I was wondering, how to do it if we need to compute modulo $$$p^k$$$ rather than just $$$p$$$. Turns out one can do it efficiently if $$$p$$$ is small. Great thanks to Endagorion for sharing it with me!

Task formulation and outline of result

To clarify the task a bit, our ultimate goal here is to be able to compute e.g. binomial coefficients modulo $$$p^k$$$. Thus, what we really need is to represent $$$n! = p^t a$$$, where $$$\gcd(a, p)=1$$$, and then report $$$t$$$ and $$$a$$$ modulo $$$p^k$$$. This is sufficient to also compute $$$\binom{n}{r}$$$ modulo $$$p^k$$$.

We will show that, assuming that polynomial multiplication of size $$$n$$$ requires $$$O(n \log n)$$$ operations, and assuming that arithmetic operations modulo $$$p^k$$$ take $$$O(1)$$$, we can find $$$t$$$ and $$$a$$$ in $$$O(d^2 + dk\log k)$$$, where $$$d=\log_p n$$$. It requires $$$O(pk\log^2 k)$$$ pre-computation that takes $$$O(pk \log k)$$$ memory.

Motivational example

Xiaoxu Guo Contest 3 — Binomial Coefficient. Given $$$1 \leq n,k \leq 10^{18}$$$, find $$$\binom{n}{k}$$$ modulo $$$2^{32}$$$.

There are also some projecteuler problems that require it, including with several queries of distinct $$$n$$$ and $$$k$$$.

Algorithm with polynomials

Let $$$k=\lfloor \frac{n}{p} \rfloor$$$, then we can represent $$$n!$$$ as

So, there is a part contributed by numbers divisible by $$$p$$$, which can be reduced to computing $$$k!$$$, and a part contributed by numbers not divisible by $$$p$$$. Let $$$n = a_0 + a_1 p + \dots + a_d p^d$$$, then the later can be represented as follows:

To compute it quickly, we can define a family of polynomial $$$P_{i,a}(x)$$$ for $$$0 \leq a \leq p$$$ such that

so the value of the factorial would be represented as

and expanding $$$k!$$$ it then rewrites into

which simplifies as

Now, what would it take us to use this setup? First of all, notice that $$$P_{i,a}(x)$$$ can be computed from one another:

Note that for shorter and more consistent implementation, this recurrent formula also mostly works for $$$i=1$$$ if we put $$$P_{0, p}(x) = x+1$$$, but for $$$P_{1,p}$$$ we should go up to $$$p-2$$$ instead of $$$p-1$$$. We should also note that for $$$P_{i,a}(x)$$$, we only care for coefficients up to $$$x^{\lfloor x/i \rfloor}$$$, as the larger ones are divisible by $$$p^k$$$.

This allows to compute $$$P_{i,a}(x)$$$ in $$$O(\frac{pk \log k}{i})$$$ for all $$$a$$$ for a given $$$i$$$. Over all $$$i$$$ from $$$1$$$ to $$$k$$$ it sums up to $$$O(pk \log^2 k)$$$ time and $$$O(p k \log k)$$$ memory. Then, evaluating all the polynomials for any specific $$$a$$$ requires $$$O(d^2 + dk \log k)$$$ operations, where $$$d = \log_p n$$$.

As it requires some manipulations with polynomials, I implemented it in Python with sympy just as a proof of concept:

from sympy import poly
from sympy.abc import x

p = 2
k = 6
mod = p**k

# replace P(x) with P(arg)
def subs_expand(P, arg):
    return poly(P.subs(x, arg).as_expr().expand(), x, modulus=mod)

P = [[poly(x+1, modulus=mod)]]

# P[i] is constant for larger i
for i in range(1, k+2):
    P.append([poly(1, x, modulus=mod)])
    for j in range(p-(i==1)):
        P[-1].append(P[-1][-1]*subs_expand(P[-2][-1], p*x+j) % poly(x**(k//i+1), modulus=mod))

def fct(n, h=0):
    if n == 0:
        return 0, 1
    t = n // p
    d = n % p
    k, a = fct(t, h+1)
    k += t
    for i in range(h+1):
        a = a * P[min(len(P)-1, i+1)][d].subs(x, t) % mod
    return k, a

Algorithm with $$$p$$$-adic logarithms and exponents

The algorithm above requires some heavy machinery on polynomials. I also found out an alternative approach that allows to compute $$$t$$$ and $$$a$$$ for any given $$$n$$$ divisible by $$$p$$$ in $$$O(pk)$$$. It doesn't rely on polynomial operations, except for Lagrange interpolation.

Let $$$n=pt+b$$$, where $$$0 \leq b < p$$$, then we can represent its factorial as

We can further rewrite it as

Let's learn how to compute the product

as with it we can represent the factorial as

We can take a $$$p$$$-adic logarithm from the product to get

We can precompute sums of $$$i^{-z}$$$ for each $$$z$$$ up to $$$k$$$ and $$$b$$$ up to $$$p-1$$$ in $$$O(pk)$$$, and we can find the sum of $$$j^z$$$ for $$$j$$$ from $$$0$$$ to $$$t$$$ in $$$O(z)$$$ via Lagrange interpolation. Therefore, with $$$O(pk)$$$ precomputation, we can compute $$$\log A(b, t)$$$ in $$$O(k^2)$$$, which then allows to find $$$n!$$$ in $$$O(dk^2)$$$, where $$$d = \log_p n$$$. I implemented it in Python, as a proof of concept, but I didn't bother with faster sums over $$$i$$$ and $$$j$$$:

p = 3
k = 7
mod = p**k
phi = (p-1)*p**(k-1)

def inv(z):
    return pow(z, phi-1, mod)

def valuation(z):
    ans = 0
    while z % p == 0:
        ans += 1
        z //= p
    return ans

def sumi(z, b):
    return sum(pow(inv(i), z, mod) for i in range(1, b+1)) % mod

def sumj(z, t):
    return sum(pow(j, z, mod) for j in range(t+1)) % mod

def logA(b, t):
    ans = 0
    for z in range(1, 40):
        pw = valuation(z)
        ans += (-1)**(z+1) * p**(z-pw) * inv(z//p**pw) * sumi(z, b) * sumj(z, t) % mod
    return ans % mod

F = [1]
for i in range(1, p):
    F.append(F[-1]*i % mod)

def pexp(lg):
    assert lg % p == 0
    ans = 0
    fct = 1
    for z in range(40):
        pw = valuation(fct)
        ans += pow(p, z-pw, mod) * inv(fct//p**pw) * pow(lg//p, z, mod) % mod
        fct = fct*(z+1)
    assert ans % p == 1
    return ans % mod

def fct(n):
    if n == 0:
        return 0, 1
    t, b = n // p, n % p
    k, a = fct_test(t)
    k += t
    a = a * pow(F[-1], t, mod) * F[b] % mod
    a = a * pexp(logA(b, t) + logA(p-1, t-1) - logA(b, t-1)) % mod
    return k, a

Note: Due to inherent properties of $$$p$$$-adic logarithms and exponents, the algorithm only works with $$$p>2$$$, and it might return $$$-n!$$$ instead with $$$p=2$$$. This is because $$$\exp \log z = -z$$$ with $$$p=2$$$ when $$$z$$$ has remainder $$$3$$$ modulo $$$4$$$, and it should be tracked separately.

Some other approaches

There is also a blog by prabowo, and a scholarly article by Andrew Granville. I'm not sure whether they talk about the same algorithms as described here.