AtCoder ABC 156

February 23, 2020

ABC156 の自分用メモ

A	B	C	D	E	F
o	o	o	x	x	x

A, B, Cは問題なく解くことができたが、Dで引っかかった。

D問題

https://atcoder.jp/contests/abc156/tasks/abc156_d

aとbを考慮せずn種類の花から1本以上選んで花束を作る方法は

${}_n C _1 + {}_n C_2 + {}_n C_3 + \dots + {}_n C_n = 2^n - 1$

で求められる(あるいはn種類の花それぞれについて選ぶ選ばないを考えて $2^n$ 、すべてを選ばない場合を引いて $2^n - 1$ )。そこから ${}_n C_a$ と ${}_n C_b$ を引いた値が求める答えとなることはわかったが、 $2 \le n \le 10^9$ という条件を見て、 ${}_n C_a$ と ${}_n C_b$ が制限時間以内に計算できないと判断してしまった。解いてる最中は累乗と階乗が頭でごっちゃになっていたり、累乗etcのアルゴリズムについての理解不足が原因。以下、累乗、累乗の逆元、二項係数それぞれについてまとめる。

累乗

$3^{16}$ について求めたいとする。愚直に求めると

p = 1
for 1...16 {
    p *= 3;
}

のように　 $O(n)$ となる。しかし、

$\begin{eqnarray} 3^{16} & = & 3^8 * 3^8 \\ 3^{8} & = & 3^4 * 3^4 \\ 3^{4} & = & 3^2 * 3^2 \\ 3^{2} & = & 3 * 3 \end{eqnarray}$

と計算できることを考慮すると、

p = 3
for i = 1...4 {
    p = p * p
}

で求めることができ、この場合 $O(\log{n})$ となる。求めたい累乗は必ずしも2のn乗とはならないから一般的に $3^n$ を求めるときに、この方法を応用できないかを考える。 $3^{2n}$ は $O(\log{n})$ で計算できることがわかっているので、nを2の累乗の和に分解できれば、

$3^{19} = 3^{2^0 + 2^1 + 2^4} = 3^{2^0} * 3^{2^1} * 3^{2^4}$

のように計算できる。累乗和への分解は、nを二進数にしたとき、i桁目が1かどうかを調べることでできる。これを踏まえると、以下のコードで $3^n$ が $O(\log{n})$ で計算できる。

ans = 1
p = 3
while (n > 0) {
    if (n & 1) ans *= p
    p = p * p
    n = n >> 1
}

$a^n$ も同様に求められる。今回のD問題のようにmodをとる必要がある問題も多いが、累乗は乗算しかおこなわないので、各計算ごとにmodをとれば $a^{n}\bmod p$ も計算できる。

累乗の逆元

modを取らない場合は、上の方法で累乗を求めて割ればいいだけなので、 $\bmod p$ における $a^n$ の逆元を求める方法についてまとめる。以下pを法とする。求める方法は2つある。

Fermatの小定理を使う
拡張ユークリッド互除法を使う

どちらも計算量は $O(\log{p})$ となる。前者の場合、Fermatの小定理からaの逆元は $a^{p-2}$ となる。 $a^{p-2}$ は上の累乗を求めるアルゴリズムを用いて $O(\log{p})$ で求められる。したがって、累乗の逆元を求めたい場合は、累乗を求めて、その逆元を求めればよいので、 $O(\log{n} + \log{p})$ となる。

二項係数

${}_n C_{k} = \frac{n!}{k!(n-k)!} = \frac{n(n-1)\cdots(n-k+1)}{k!}$

今回の問題では、 $0 \le a,\,b \le 2\times 10^5$ なので、 ${}_n C_{k}$ の分母、分子はそれぞれ $O(a)$ で計算できるので間に合う。しかし、aがより大きいとき、階乗を効率的に計算する方法があるのか気になったので調べたところ $n! \bmod p$ を $O(p^{1/2}\log{p})$ で計算する方法があるらしい[2]。

参考

[1]「1000000007 で割ったあまり」の求め方を総特集！〜逆元から離散対数まで〜 https://qiita.com/drken/items/3b4fdf0a78e7a138cd9a
[2] 階乗 mod 素数 https://min-25.hatenablog.com/entry/2017/04/10/215046