小さい素数modの二項係数

Lucasの定理をPythonで実装する。

以下 $p$ を素数とする。 パスカルの三角形を使って、$p$ 未満の数についての二項係数を計算しておく。

pascal = [[0] * p for _ in range(p)]
pascal[0][0] = 1
for i in range(1, p):
    pascal[i][0] = 1
    for j in range(1, i + 1):
        pascal[i][j] = (pascal[i - 1][j - 1] + pascal[i - 1][j]) % p

$x$ を $p$ 進法で表した配列を返す関数。

def base(x):
    res = []
    y = x
    while y:
        res.append(y % p)
        y //= p
    return res

$_{x} C _ {y}$ を $p$ で割ったあまりを返す関数。

def cmb(x, y):
    if x < y or x < 0 or y < 0:
        return 0
    x_lst = base(x)
    y_lst = base(y)
    for _ in range(len(x_lst) - len(y_lst)):
        y_lst.append(0)
    res = 1
    for i, j in zip(x_lst, y_lst):
        res *= pascal[i][j]
        res %= p
    return res

使用例

ARC 117 C - Tricolor Pyramid

p = 3
# 中略
n = int(input())
dic = {"W": 0, "B": 1, "R": 2}
rev = ["W", "B", "R"]
a = [dic[x] for x in input()]
ans = 0
for i in range(n):
    ans += cmb(n - 1, i) * a[i]
    ans %= p
ans *= pow(2, n - 1, p)
ans %= p
print(rev[ans])

Submission #21876484 - AtCoder Regular Contest 117