4200000000Hzでゴリ押す！！！(1)

執筆者: 終に鮭

最終更新: 2021/05/09

プログラミング

何でもかんでもアルゴリズムとマシンパワー（主にCPU）でゴリ押します。そういう企画です。

「AtCoder水下位競プロerである私ならこういう思考をするぞ」っていうのを書きます（個人の意見ですが）。

2021年東京大学(理系) 数学問4(4)

初回は東大の入試問題をやります。この問題をパワーでゴリ押すのは $O(1)$ 番煎じになりますが、気にしないことにします。

実はこれ春休みの帰省中に数人で集まって解いた問題の一つで、僕だけパワーで解決しようとしてました。

問題

${}_{2021}C_{37}\bmod 4$ を求めよ。

本来はいろいろと前置きがあって、それを利用して引数を落として計算するんですが、そんなの私には分かりません。

さて、この問題をどうやって解いてやりましょうか。

解法1

${}_{2021}C_{37}$ を求める

で、問題は ${}_{2021}C_{37}$ がでかすぎることです。どの程度か簡単に見積もってみましょう。

$$
\begin{aligned}
{}_{2021}C_{37}
&=\frac{2021 \times \dots \times 1985}{37 \times \dots \times 1} \\
&>\left(\frac{1985}{37}\right)^{37} \\
&>2^{212.58}
\end{aligned}
$$
困りました。これでは128ビット符号付き整数(→GCCの__int128、C++ Boostのint128_t、Rustのi128など)でも収まらず、直接計算するのは厳しいです。

十分な長さの固定長整数型や多倍長整数を利用すれば出来ないこともないので一応やってみましょう。
Pythonのint型は多倍長なので何食わぬ顔で書けば実装できます。

n = 2021
r = 37
MOD = 4

ans = 1
for i in range(min(r, n - r)):
  ans = ans * (n - i) // (1 + i)

print(f'{n}C{r} = {ans}')
print(f'{n}C{r} % {MOD} = {ans % MOD}')

実行結果は次のようになりました。

2021C37 = 10549453167137272405699426118445044729600037976073052501105621045645733312800275
2021C37 % 4 = 3

な、成程……！
$${}_{2021}C_{37}=10549453167137272405699426118445044729600037976073052501105621045645733312800275$$
なんだ……！そして $10549453167137272405699426118445044729600037976073052501105621045645733312800275$ を $4$ で割ったあまりは $3$ なんだ……！

多倍長整数での $N \times M$ および $N / M$ $(N\gg M)$ が $O(\log N \log\log N)$ で計算できる（例えばFFTとNewton-Raphson法で実装）という仮定のもとで、
この方法による ${}_{n}C_{r} (r \leq n/2)$ の計算量は大きめに見積もると、
$$
\begin{aligned}&
O(\log {}_{n}C_{1} \log\log {}_{n}C_{1})
+ \dots + O(\log {}_{n}C_{r} \log\log {}_{n}C_{r}) \\
=& O(\log n \log\log n) + \dots + O(\log n^r \log\log n^r) \\
=& O(r^2 \log n \log\log n)
\end{aligned}
$$
となります（この評価は割とガバガバですが）。

解法2

${}_{n}C_{r}=\dfrac{{}_{n}P_{r}}{r!}$ とオイラーの定理を利用する

${}_{2021}C_{37} \bmod 4$ を計算しようと思ったらこの方法が一番に思いつきました。

オイラーの定理は、競技プログラミング界隈では常識であるところのフェルマーの小定理の一般化です。
この定理は次のようなものです。

$m$ 以下の正の整数のうち $m$ に互いに素であるものの個数を $\varphi(m)$ とする（言い換えれば、$\varphi(m)=|\{n ; \gcd(m,n)=1, 1 \leq n \leq m\}|$）。
$a$ と $m$ が互いに素であるとき、
$$a^{\varphi(m)} \equiv 1 \pmod{m}$$

これの何が嬉しいって、分子 $a$ と分母 $b$ を $\bmod m$ で別々に計算できるので、オーバーフローさせずに固定長整数型で計算できます。

そして分子に $\bmod m$ における分母の逆元 $b^{-1}$ を掛ける事によって求めることが出来ます。できるはずでした。

しかしよーく考えてみましょう。${}_{2021}P_{37},37! \pmod{4}$ っていくつになるでしょうか？

……そうです。$0$ になります。$0$ に何掛けたって $4$ で割ったあまりが $1$ になることはありません。
すなわち、$b \equiv 0$ の逆元 $b^{-1}$ なんてものは存在しないので、この方法は使えません。
${}_{2021}C_{37} \bmod (10^9+7)$ であれば求められるんですけどね……

解法2'

ただこれを回避する方法があります。

$${}_{n}C_{r}=\frac{a}{b},a={}_{n}P_{r},b=r!$$
とします。ここで、
$$
\begin{aligned}
a &=2^{\alpha_0} \times 3^{\alpha_1}
\times 5^{\alpha_2} \times 7^{\alpha_3} \times \dots \\
b &=2^{\beta_0} \times 3^{\beta_1}
\times 5^{\beta_2} \times 7^{\beta_3} \times \dots
\end{aligned}
$$
のように素因数分解し、約分することで
$$
{}_{n}C_{r}=\frac{a}{b}
=2^{\alpha_0-\beta_0} \times 3^{\alpha_1-\beta_1}
\times 5^{\alpha_2-\beta_2} \times 7^{\alpha_3-\beta_3} \times \dots
$$
として計算することができます（${}_{n}C_{r}$ は整数であるから、任意の $i$ について $\alpha_i \geq \beta_i$ が成り立つ）。

実際にこれを実装してみました。

from collections import defaultdict

n = 2021
r = 37
MOD = 4

spf = list(range(n + 1))
for i in range(2, n + 1):
  if spf[i] != i:
    continue
  for j in range(i * 2, n + 1, i):
    if spf[j] == j:
      spf[j] = i

pf = defaultdict(lambda:0)
for i in range(min(r, n - r)):
  m = n - i
  while m > 1:
    pf[spf[m]] += 1
    m //= spf[m]
  m = 1 + i
  while m > 1:
    pf[spf[m]] -= 1
    m //= spf[m]

ans = 1
for p, v in pf.items():
  ans *= pow(p, v, MOD)
  ans %= MOD

print(f'{n}C{r} % {MOD} = {ans}')

結果は当然解法1と同じく3です。

各自然数を $O(\sqrt{n})$ で素因数分解して全体 $O(r \sqrt{n})$ で求める方法もありますが、$r \fallingdotseq n/2$ の場合も想定して、エラトステネスの篩と同様の前計算 $O(n \log\log n)$ でSPF（Smallest Prime Factor）配列を求め、一つの自然数の素因数分解は $O(\log n)$ で、全体 $O(n \log\log n + r \log n)$ でできるようにしました。

全体の計算量は、$O(n \log\log n + r \log n)$ です。
東大の問題の場合 $r$ が小さいので微妙ですが、$r \fallingdotseq n/2$ の場合はこちらのほうが早そうです。

解法3

パスカルの三角形を使う

今回の大本命です。

解法2では計算途中で $4 \equiv 0$ で掛けたり割ったりするから計算できなかったのが悪いので、いっそのこと足し算だけで求めてしまおうという考え。

$$
{}_{n}C_{r}=\begin{cases}
0 &(n<r) \\
1 &(r=0) \\
{}_{n-1}C_{r} + {}_{n-1}C_{r-1} &(\text{otherwise})
\end{cases}
$$
この漸化式を使えば解けます。軽く証明してみましょう。

$$
\begin{aligned}
{}_{n-1}C_{r} + {}_{n-1}C_{r-1}
&=\frac{(n-1)!}{(n-r-1)!r!}+\frac{(n-1)!}{(n-r)!(r-1)!} \\
&=\frac{n-r}{n-r}\frac{(n-1)!}{(n-r-1)!r!}
+\frac{r}{r}\frac{(n-1)!}{(n-r)!(r-1)!} \\
&=(n-r)\frac{(n-1)!}{(n-r)!r!}+r\frac{(n-1)!}{(n-r)!r!} \\
&=\frac{n!}{(n-r)!r!}={}_{n}C_{r} \\
\end{aligned}
$$

ではこれを実装します。

実装1: 再帰関数

漸化式と言えば再帰ですよね。見たまんま実装できるので良いと思います。

n = 2021
r = 37
MOD = 4

def comb(n, r, mod):
  if n < r:
    return 0
  elif r == 0:
    return 1
  else:
    return (comb(n - 1, r, mod) + comb(n - 1, r - 1, mod)) % mod

print(f'{n}C{r} % {MOD} = {comb(n, r, MOD)}')

でもちょっと待ってください。これ、comb関数の呼び出し回数どうなると思います？

$$
\operatorname{call}(n,r)=\begin{cases}
1 &(n<r) \\
1 &(r=0) \\
\operatorname{call}(n-1,r) + \operatorname{call}(n-1,r-1) + 1 &(\text{otherwise})
\end{cases}
$$

として定義した $\operatorname{call}(n, r)$ 回になるんですよ。つまり少なくとも ${}_{n}C_{r}$ より大きい回数呼び出されます。
大きく見積もれば $O(2^n)$。$n=2021$ で、実行していい計算量ではありません。

実装2: メモ化再帰

上の実装の問題点は、${}_{n}C_{r}$ を計算するために、${}_{n-x}C_{r-y}$ の計算を ${}_{x}C_{y}$ 回もしてしまっているということです。

ここで、メモ化という手法を取ることによって、各 $(x,y)$ に対して ${}_{n-x}C_{r-y}$ を $1$ 回のみ計算すれば良いようにすることができます。

from sys import setrecursionlimit
setrecursionlimit(3000)

n = 2021
r = 37
MOD = 4

memo = [[None] * (r + 1) for _ in range(n + 1)]
def comb(n, r, mod):
  if n < r:
    return 0
  elif r == 0:
    return 1
  elif memo[n][r] != None:
    return memo[n][r]
  else:
    memo[n][r] = (comb(n - 1, r, mod) + comb(n - 1, r - 1, mod)) % mod
    return memo[n][r]

print(f'{n}C{r} % {MOD} = {comb(n, r, MOD)}')

コードを見ればそう難しい話ではないですね。単に二次元配列に計算済みの値を保存して、計算済みならば保存した値を返すようにしているだけです。

計算量は $O(nr)$ ですが、空間計算量も $O(nr)$ なのが痛いところです。

実装3: 動的計画法

動的計画法とは、問題を部分問題に分割し、自明な部分問題から遷移していって元の問題を解く手法です。
メモ化再帰もこれに含まれますが、単純再帰は含まれません（動的計画法の要件に計算結果の記録が含まれるため）。

ここでは、ボトムアップに（すなわち、自明な部分問題から順に）求める方法を指します。

パスカルの三角形の $0$ 行目から $n$ 行目まで順番に埋めていくイメージですね。

n = 2021
r = 37
MOD = 4

dp = [[0] * (r + 1) for _ in range(n + 1)]
dp[0][0] = 1

for i in range(1, n + 1):
  dp[i][0] = 1
  for j in range(1, r + 1):
    dp[i][j] = (dp[i - 1][j] + dp[i - 1][j - 1]) % MOD

print(f'{n}C{r} % {MOD} = {dp[n][r]}')

これもやはり $O(nr)$ です。

実装3': SIMD

普通に順次計算すると1クロックで1命令で、計算が終わってまた1命令出すんですが、その1命令に複数のオペランドに対する同一の演算を詰め込むことができます。これがSIMDです。
ベクトルの足し算や内積などが一番よく利用されている例だと思います。

さて、PythonでもSIMDをすることができます。
Numpyというライブラリを利用することで、配列同士の和を取ったりでき、自動でSIMDを使うようになります。内部的に PADDB 命令とかを出してるんでしょうね。

import numpy as np

n = 2021
r = 37
MOD = 4

dp = np.zeros((n + 1, r + 1), np.int8)
dp[0, 0] = 1

for i in range(1, n + 1):
  dp[i, 0] = 1
  dp[i, 1:] = (dp[i - 1, 1:r + 1] + dp[i - 1, 0:r]) % MOD

print(f'{n}C{r} % {MOD} = {dp[n, r]}')