Learning Algorithms

アルゴリズムの勉強メモ

Inverse of factorialsをO(N)で列挙する

Inverse of factorialsをO(N)で列挙する

はじめに

素数 $p$ を法として、 $n$ 以下の階乗の逆元を $O(n)$ で列挙しちゃおうというアレを知っていますか?(僕は知りませんでした。。。)

(コードが読みにくくなる割には)そんなに高速になるわけではないためか、実際に使っている人はほとんど見たことがないので、書いておこうと思います。

まず典型的な計算量 $O(n \log n)$ での列挙は次のようなものでした。ここでは $10^9 + 7$ (素数)を法として、$5000000$ 以下の階乗とその逆元を列挙しています。

const int MOD = 1e9 + 7;

const int N = 5000001;
long long fact[N];
long long invfact[N];

long long Inv(long long a) {
        long long res = 1;
        long long n = MOD - 2;
        while (n > 0) {
                if (n & 1) res = res * a % MOD;
                a = a * a % MOD;
                n >>= 1;
        }
        return res;
}

void init() {
        fact[0] = fact[1] = 1;
        for (int i = 2; i < N; i ++) fact[i] = fact[i - 1] * i % MOD;
        invfact[0] = invfact[1] = 1;
        for (int i = 2; i < N; i ++) invfact[i] = invfact[i - 1] * Inv(i) % MOD;
}

フェルマーの小定理と繰り返し二乗法を使って各数の逆元をそれぞれ $O(log n) $で求めています。いえい。

本題

ところで、上の計算に無駄があるとすれば、それは毎回Inv(i)を呼び出して逆元を計算していることです。そこで、Inv[i]をあらかじめ $O(n)$ で求めておく方法を考えます。

まず、次の式を眺めることにします。これは $p$ を $i$ で割ったときの計算式そのものになっています。つまり、 $ \frac {p}{i}$ が商、 $p\ \%\ i$ がその余りになっています。

$p = i * \frac {p}{i} + p\ \%\ i$

この式を $p$ を法として考えると、

$0 \equiv i * \frac {p}{i} + p\ \%\ i$

移項して両辺に $(\frac{p}{i})^{-1}$ をかけると

$i \equiv -\ (\frac {p}{i})^{-1} * (p\ \%\ i)$

両辺の逆元をとると

$i^{-1} \equiv -\ (\frac {p}{i})(p\ \%\ i) ^ {-1}$

さて、ここで注目したいのは、$p\ \%\ i < i$ であるため、小さいものから順に逆元を計算していけるということです。

はい、これで $O(n)$ で逆元を前計算しておけることがわかったので、全体でも計算量は $O(n)$ になりました。これをそのまま実装すると次のようになります。ただし負数 $-x$ は 正数 $p-x$ として扱っています。

const int MOD = 1e9 + 7;

const int N = 5000001;
long long fact[N];
long long invfact[N];
long long inv[N];

void init() {
        fact[0] = fact[1] = 1;
        for (int i = 2; i < N; i ++) fact[i] = fact[i - 1] * i % MOD;
        inv[1] = 1;
        for (int i = 2; i < N; i ++) inv[i] = (MOD - MOD / i) * inv[MOD % i] % MOD;
        invfact[0] = invfact[1] = 1;
        for (int i = 2; i < N; i ++) invfact[i] = invfact[i - 1] * inv[i] % MOD;
}

簡単なベンチマークをしたところ、およそ3倍程度高速に計算されている感じでした。

この実装の問題点を挙げるとすれば、なぜこれで逆元が計算されるのかがぱっとわからないことくらいでしょうか。プログラム中の他の部分で逆元を多用する場合にはこのように前計算しておくのもありかもしれません。