The goal is to find \$f(n,r)\mod(m)\$ for the given values of \$n\$, \$r\$, \$m\$, where \$m\$ is a prime number.

$$f(n,r) = \dfrac{F(n)}{F(n-r) \cdot F(r)}$$


$$F(n) = 1^1 \cdot 2^2 \cdot 3^3 \cdot \ldots \cdot n^n$$

Here is the Python code snippet I've written:

    n_r = n - r 
    num = 1
    den = 1

    if n_r > r:
        for j in xrange(n_r + 1, n + 1): 
            num = num*(j**j)    
        for k in xrange(2, r + 1):
            den = den*(k**k)

        answer = num/den
        ans = answer%m
        print ans

        for j in xrange(r + 1, n + 1): 
            num = num*(j**j)    
        for k in xrange(2, n_r + 1):
            den = den*(k**k)

        answer = num/den
        ans = answer%m
        print ans    

This code runs for small values of \$n\$ (when \$n <= 100\$). But for large values of \$n\$ ~ \$10^5\$, the code seems to be inefficient, exceeding the time limit of 2 sec. How can I generalize and optimize this computation?


2 Answers 2


Your function is slow for large values of n because the intermediate values of num and den quickly grow to huge integers, making the multiplication slow.

You can improve that by reducing each intermediate result "modulo m", so that all numbers will always be in the range 0 ... m-1.

The final division num/den must then be computed as a "modular division":

def mod_div(a, b, m):
    """modular division

    Returns a solution c of a = b * c mod m, or None if no such number c exists.
    g, x, y = egcd(b, m)
    if a % g == 0:
        return (a * (x // g)) % m
        return None

using the "Extended Euclidean Algorithm, for example with the code from http://en.wikibooks.org/wiki/Algorithm_Implementation/Mathematics/Extended_Euclidean_algorithm#Recursive_algorithm

def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)

The if/else can be avoided, either as @Janos suggested, or by replacing r by n-r if the latter is smaller. Your function would then look like this:

def f(n, r, m):

    if n - r < r:
        r = n - r

    num = 1
    den = 1

    for j in xrange(n - r + 1, n + 1): 
        num = (num * (j ** j % m)) % m
    for k in xrange(2, r + 1):
        den = (den * (k ** k % m)) % m

    return mod_div(num, den, m)

At least for large values of n this should be faster. I made a test with f(10000, 100, 747164718467):

  • Your code: 25 seconds
  • Above code: 0.2 seconds

A possible further improvement would be to compute the powers j ** j % m also with a "modular power" method, such as power_mod() from http://userpages.umbc.edu/~rcampbel/Computers/Python/lib/numbthy.py. The reduced the time for calculating f(10000, 100, 747164718467) to 0.05 seconds on my computer.

  • \$\begingroup\$ The proposed code does not handle the "modular inverse does not exist" case. For example f(6,2,2) makes den == 0. \$\endgroup\$ Nov 16, 2014 at 19:02
  • 1
    \$\begingroup\$ @JanneKarila: You are right. Actually the first step would be to prove that f(n,r) is an integer at all (which is implied in the question). I think that can be done similar as here for binomial coefficients. And then I think that if the denominator has the prime factor m then the numerator has at least the factor m^2, so the modulo result would be zero in that case. \$\endgroup\$
    – Martin R
    Nov 16, 2014 at 19:58
  • \$\begingroup\$ @JanneKarila: I have updated the code so that it works for those cases as well. \$\endgroup\$
    – Martin R
    Nov 17, 2014 at 21:00

The logic in the two branches of the if statement is essentially the same, the only difference is in the parameters of xrange. You can generalize the logic using a variable x, set to max(n_r, r):

x = n_r if n_r > r else r

for j in range(x + 1, n + 1):
    num *= j ** j
for k in range(2, n - x + 1):
    denom *= k ** k

answer = num / denom
ans = answer % m

I made some other changes too:

  • For compatibility with Python 3:
    • Use print(...) instead of print ...
    • Use range instead of xrange
  • More generous spacing around operators, for PEP8
  • Use p *= q instead of p = p * q
  • Renamed den to denom

As for the speed issue... That seems more of a math problem than Python, and I don't know the answer for that.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.