8

I need to calculate nCr mod p efficiently. Right now, I have written this piece of code, but it exceeds the time limit. Please suggest a more optimal solution.

For my case, p = 10^9 + 7 and 1 ≤ n ≤ 100000000

I have to also make sure that there is no overflow as nCr mod p is guaranteed to fit in 32 bit integer, however n! may exceed the limit.

def nCr(n,k):
    r = min(n-k,k)
    k = max(n-k,k)
    res = 1
    mod = 10**9 + 7

    for i in range(k+1,n+1):
        res = res * i
        if res > mod:
            res = res % mod

    res = res % mod
    for i in range(1,r+1):
        res = res/i
    return res

PS : Also I think my code may not be completely correct. However, it seems to work for small n correctly. If its wrong, please point it out !

OneMoreError
  • 7,518
  • 20
  • 73
  • 112

2 Answers2

12

From http://apps.topcoder.com/wiki/display/tc/SRM+467 :

long modPow(long a, long x, long p) {
    //calculates a^x mod p in logarithmic time.
    long res = 1;
    while(x > 0) {
        if( x % 2 != 0) {
            res = (res * a) % p;
        }
        a = (a * a) % p;
        x /= 2;
    }
    return res;
}

long modInverse(long a, long p) {
    //calculates the modular multiplicative of a mod m.
    //(assuming p is prime).
    return modPow(a, p-2, p);
}
long modBinomial(long n, long k, long p) {
// calculates C(n,k) mod p (assuming p is prime).

    long numerator = 1; // n * (n-1) * ... * (n-k+1)
    for (int i=0; i<k; i++) {
        numerator = (numerator * (n-i) ) % p;
    }

    long denominator = 1; // k!
    for (int i=1; i<=k; i++) {
        denominator = (denominator * i) % p;
    }

    // numerator / denominator mod p.
    return ( numerator* modInverse(denominator,p) ) % p;
}

Notice that we use modpow(a, p-2, p) to compute the mod inverse. This is in accordance to Fermat's Little Theorem which states that (a^(p-1) is congruent to 1 modulo p) where p is prime. It thus implies that (a^(p-2) is congruent to a^(-1) modulo p).

C++ to Python conversion should be easy :)

batbaatar
  • 5,448
  • 2
  • 20
  • 26
2

About the last question: I think that the mistake in your code is to compute the product, reduce it modulo k, and then divide the result by r!. This is not the same as dividing before reducing modulo k. For example, 3*4 / 2 (mod 10) != 3*4 (mod 10) / 2.

Armin Rigo
  • 12,048
  • 37
  • 48