Evaluating f(N,R) mod(m)

Question

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)}$$

where

$$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

    else:
        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?

Answer 1

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
    else:
        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)
    else:
        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.

Answer 2

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
print(ans)

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.