Search code examples
c++algorithmmodulobinomial-coefficients

nCk modulo p when n % p or k % p == 0


I'm trying to solve a coding challenge on hacker rank which requires one to calculate binomial coefficients mod a prime, i.e.

nchoosek(n, k, p) 

I'm using the code from this answer that works for the first three sets of inputs but begins failing on the 4th. I stepped through it in the debugger and determined that the issue arises when:

n % p == 0 || k % p == 0

I just need to know how to modify my current solution to handle the specific cases where n % p == 0 or k % p == 0. None of the answers I've found on stack exchange seem to address this specific case. Here's my code:

#include <iostream>
#include <fstream>

long long FactorialExponent(long long n, long long p)
{
    long long ex = 0;
    do
    {
        n /= p;
        ex += n;
    }while(n > 0);
    return ex;
}

unsigned long long ModularMultiply(unsigned long long a, unsigned long long b, unsigned long p) {
  unsigned long long a1 = (a >> 21), a2 = a & ((1ull << 21) - 1);
  unsigned long long temp = (a1 * b) % p;   // doesn't overflow under the assumptions
  temp = (temp << 21) % p;                  // this neither
  temp += (a2 * b) % p;                       // nor this
  return temp % p;
}

unsigned long long ModularInverse(unsigned long long k, unsigned long m) {
    if (m == 0) return (k == 1 || k == -1) ? k : 0;
    if (m < 0) m = -m;
    k %= m;
    if (k < 0) k += m;
    int neg = 1;
    unsigned long long p1 = 1, p2 = 0, k1 = k, m1 = m, q, r, temp;
    while(k1 > 0) {
      q = m1 / k1;
      r = m1 % k1;
      temp = q*p1 + p2;
      p2 = p1;
      p1 = temp;
      m1 = k1;
      k1 = r;
      neg = !neg;
    }
    return neg ? m - p2 : p2;
}

// Preconditions: 0 <= k <= min(n,p-1); p > 1 prime
unsigned long long ChooseModTwo(unsigned long long n, unsigned long long k, unsigned long p)
{
    // reduce n modulo p
    n %= p;

    // Trivial checks
    if (n < k) {

        return 0;
    }

    if (k == 0 || k == n) {
        return 1;
    }

    // Now 0 < k < n, save a bit of work if k > n/2
    if (k > n/2) {
        k = n-k;
    }
    // calculate numerator and denominator modulo p
    unsigned long long num = n, den = 1;
    for(n = n-1; k > 1; --n, --k)
    {
        num = ModularMultiply(num, n, p);
        den = ModularMultiply(den, k, p);
    }

    den = ModularInverse(den,p);
    return ModularMultiply(num, den, p);
}

// Preconditions: 0 <= k <= n; p > 1 prime
long long ChooseModOne(long long n, long long k, const unsigned long p)
{
    // For small k, no recursion is necessary
    if (k < p) return ChooseModTwo(n,k,p);
    unsigned long long q_n, r_n, q_k, r_k, choose;
    q_n = n / p;
    r_n = n % p;
    q_k = k / p;
    r_k = k % p;

    choose = ChooseModTwo(r_n, r_k, p);
    // If the exponent of p in choose(n,k) isn't determined to be 0
    // before the calculation gets serious, short-cut here:
    // if (choose == 0) return 0;
    return ModularMultiply(choose, ChooseModOne(q_n, q_k, p), p);

}

unsigned long long ModularBinomialCoefficient(unsigned long long n, unsigned long long k, const unsigned long p)
{
    // We deal with the trivial cases first
    if (k < 0 || n < k) return 0;
    if (k == 0 || k == n) return 1;
    // Now check whether choose(n,k) is divisible by p
    if (FactorialExponent(n, p) > FactorialExponent(k, p) + FactorialExponent(n - k, p)) return 0;
    // If it's not divisible, do the generic work
    return ChooseModOne(n, k, p);
}

int main() {

  //std::ifstream fin ("input03.txt");
  std::ifstream fin ("test.in");

  int kMod = 1000003;
  int T;
  fin >> T;
  int N = T;
  //std::cin >> T;

  unsigned long long n, k;
  unsigned long long a, b;
  int result[N];
  int index = 0;

  while (T--) {

    fin >> n >> k;

    a = ModularBinomialCoefficient(n - 3, k, kMod);
    b = ModularBinomialCoefficient(n + k, n - 1, kMod);

    // (1 / (n + k) * nCk(n - 3, k) * nCk(n + k, n - 1)) % 1000003
    unsigned long long x = ModularMultiply(a, b, kMod);
    unsigned long long y = ModularMultiply(x, ModularInverse((n + k), kMod), kMod);

    result[index] = y;
    index++;
  }

  for(int i = 0; i < N; i++) {
    std::cout << result[i] << "\n";
  }

  return 0;
}

Input:
6
90 13
65434244 16341234
23424244 12341234
424175 341198
7452123  23472
56000168 16000048

Output:
815483
715724
92308
903465
241972
0 <-- Incorrect, should be: 803478

Constraints:
4 <= N <= 10^9
1 <= K <= N


Solution

  • You can use Lucas' theorem to reduce the problem to ceil(log_P(N)) subproblems with k, n < p: Write n = n_m * p^m + ... + n_0 and k = k_m * p^m + ... + k_0 in base p (n_i, k_i < p are the digits), then we have

    C(n,k) = PROD(i = 0 to m, C(n_i, k_i))  (mod p)
    

    The subproblems are easy to solve, because every factor of k! has an inverse modulo p. You get an algorithm with runtime complexity O(p log(n)), which is better than that of Ivaylo's code in case of p << n, if I understand it correctly.

    int powmod(int x, int e, int p) {
        if (e == 0) return 1;
        if (e & 1) return (long long)x * powmod(x, e - 1, p) % p;
        long long rt = powmod(x, e / 2, p);
        return rt * rt % p;
    }
    
    int binom_coeff_mod_prime(int n, int k, int p) {
        long long res = 1;
        while (n || k) {
            int N = n % p, K = k % p;
            for (int i = N - K + 1; i <= N; ++i)
                res = res * i % p;
            for (int i = 1; i <= K; ++i)
                res = res * powmod(i, p - 2, p) % p;
            n /= p;
            k /= p;
        }
        return res;
    }