算法题 卢卡斯定理求组合数（Python）

题目

给定n组询问，每组询问给定三个整数a,b,p，其中pp是质数，请你输出Cba mod p的值。

输入格式

第一行包含整数n。

接下来n行，每行包含一组a,b,p。

输出格式

共n行，每行输出一个询问的解。

数据范围

1≤n≤20
1≤b≤a≤10^18
1≤p≤10^5

输入样例：

3
5 3 7
3 1 5
6 4 13

输出样例：

3
3
2

代码

# 快速幂
def qmi(a, b, p):
    res = 1
    while b:
        if b & 1:
            res = res * a % p
        a = a * a % p
        b >>= 1
    return res

# 逆元计算组合数
def C(a, b, p):
    j = a
    res = 1
    for i in range(1, b+1):
        res = j * res % p
        res = res * qmi(i, p - 2, p) % p
        j -= 1
    return res

# 卢卡斯定理
def lucas(a, b, p):
    if a < p and b < p: return C(a, b, p)
    else: return C(a % p, b % p, p) * lucas(a // p, b // p, p) % p

n = int(input())
for _ in range(n):
    a, b, p = map(int, input().split())
    print(lucas(a, b, p))