Choose K Modulo A Prime
Problem
Count the ways to pick k items out of n distinct items when the order of picking does not matter, and report the count modulo 1,000,000,007, which is prime. The value of n can reach a million, so the exact count has hundreds of thousands of digits and must never be built. Picking more items than exist gives zero ways.
Examples
Input: n = 5, k = 2
Output: 10Input: n = 1000, k = 500
Output: 159835829
Why: the exact count has about 300 digits; only its remainder is returnedInput: n = 3, k = 5
Output: 0
Why: edge case, there is no way to pick five items out of threeHints
0 / 3
The count is n factorial divided by k factorial and by (n minus k) factorial. Multiplication behaves well under a modulus, but division does not.
Under a prime modulus, dividing by a value is the same as multiplying by its inverse, and for a prime p the inverse of x is x raised to the power p minus 2.
Multiply n times n minus 1 and so on, k factors in total, reducing after every step, and separately multiply 1 up to k. Then multiply the first product by the second raised to the power of the modulus minus 2, using fast modular exponentiation. Use the smaller of k and n minus k to keep the loop short.
Solution
The count equals the falling product n times (n minus 1) down to (n minus k plus 1), divided by k factorial. Both products are easy to keep reduced modulo the prime, and the division becomes a multiplication by the modular inverse of k factorial, which Fermat's little theorem gives as that value to the power p minus 2. Because n stays below the prime, k factorial is never a multiple of it, so the inverse always exists. Swapping k for n minus k when that is smaller halves the work. Time is O(min(k, n minus k) plus log p) and space is O(1).
MOD = 1_000_000_007
def choose(n, k):
if k < 0 or k > n:
return 0 # cannot pick more than exist
k = min(k, n - k) # same count, shorter loop
top = bottom = 1
for i in range(k):
top = top * (n - i) % MOD # n (n-1) ... (n-k+1)
bottom = bottom * (i + 1) % MOD # k!
return top * pow(bottom, MOD - 2, MOD) % MOD # divide via Fermat's inverse
print(choose(5, 2)) # -> 10
print(choose(1000, 500)) # -> 159835829
print(choose(3, 5)) # -> 0
print(choose(7, 0)) # -> 1Stuck on the idea rather than the code? Permutations vs Combinations covers it.