競プロ初心者が絶対にハマる整数問題の罠

こんにちは、bitFlyerのAdvent Calendar8日目の記事です。
最近アルゴリズムの記事ばかり書いていますが、Androidエンジニアをしています。

五年前くらいにPythonでアルゴリズムの勉強をしたときに整数問題を解いていて意外なところで罠にハマったので、tipsとして紹介します。

対象読者

アルゴリズム初心者（Atcoderでいう灰〜茶）

書くこと

• 素数判定
• 最大公約数の算出
• n番目のフィボナッチ数列

これら三つについてです。
どれも一見簡単そうなのですが、素直に実装するとTLEになってしまいます、罠です。

素数判定

皆さんは素数判定をコードで書いたことはありますでしょうか。
私はなぜか以前これにハマって、Pythonを使って一行で素数判定をするコードを書いたことがあります（若干変人っぽくなってしまうので言い訳をしておくと、素数自体は好きではないです。というかむしろ2の倍数が計算しやすくて好きです）

私は素数の定義を「2 から 自分の数-1 に至るまで、どの整数でも割り切れない」と思っていて、その通りに実装しました。

val Long.isPrime
    get():Boolean {
        if (this < 2L) return false

        var i = 2L
        while (i * i <= this) {
            if (this % i == 0L) {
                return false
            }
            i++
        }
        return true
    }

こうですね。
一応、判定自体は正しくできます。
しかし、強烈に遅いです。
整数Nを判定するときにO(N)のオーダーの計算量になります。

当然のTLEです。

はい、実はもっと早く判定できる方法があって2から√Nまで割り切れなければ素数です。

val Long.isPrime
    get():Boolean {
        if (this < 2L) return false

        for (i in 2L..floor(sqrt(this.toDouble())).toLong()) {
            if (this % i == 0L) {
                return false
            }
        }
        return true
    }

これでオッケーです。（ルート部分が強烈に汚くてすみません。doubleに変換してsqrt()を使ったときに丸め誤差が怖いのですが、切り捨てで10万くらいまでは正しく素数判定できていました）

計算量はO(√N)になりました。

計算量の感覚がわかっていれば前者ではダメなのに気付けますが、初心者は一度はこれをやらかすと思ってます。

最大公約数の算出

最大公約数の定義を思い出してみます。
aとbの最大の共通因数でしたね。
例えば40と16の最大公約数は
40 = 2*2*2*5
16 = 2*2*2*2

8ですね。

「それなら、40と16の因数を算出して、共通因数を取り出せば良いのでは？？」とゴチャゴチャと実装しました。

当然のTLEです・・・

これも実はスマートな解法があります。

ユークリッドの互除法を使います（「マスターオブ整数であったな」と懐かしい気分になった記憶があります）
下記が実装です。

fun Long.gcm(other: Long): Long {
    val maxV = max(this, other)
    val minV = min(this, other)

    return when (minV) {
        0L -> maxV
        else -> minV.gcm(maxV % minV)
    }
}

上記の実装でやっていることを具体的に書くと
例えば、30と16の最大公約数を求めるとすると

30 / 16 = 1 余り 14
16 / 14 = 1 余り 2
14 / 2 = 7 余り 0

こんな要領で割る数を余りで割るという操作を繰り返して、余りが0になるまで繰り返して、その時の割る数が最大公約数にあたるので、それを返しています。

直感的にわかりそうでわからなかったので証明を読みましたが、証明は簡単なので一度目を通しても良いかもしれません

n番目のフィボナッチ数列

これも曲者ですね・・・

フィボナッチ数列とは一個前の数と二個前の数を足したら現在の数になるというものです。

a_1 = 1
a_2 = 1
は定義として決まっていて、
a_n+2 = a_n+1 + a_n
と定義される数列です。

例として
a_3 = a_1 + a_2 = 2
a_4 = a_2 + a_3 = 3

これは再帰関数を使えばすごく綺麗に書けます。

fun fibonacci(n: Long): Long {
    if (n == 1L) return 1L
    if (n == 2L) return 1L

    return fibonacci(n - 2) + fibonacci(n - 1)
}

初めて実装したとき天才であるかと勘違いしてしまうくらいスッキリしたのですが

当然のTLE・・・・・・・

フィボナッチ数列で再帰関数を学んだ人を周りでよく見かけますし、再帰関数を学ぶ上では良い題材なのだと思います。
これがなぜ遅いかというと
メソッド内でfibonacci(n-2)とfibonacci(n-1)を呼んでいますが、fibonacci(n-1)が走ったら
次はfibonacci(n-3)とfibonacci(n-2)が呼ばれます。

fibonacci(n-2)が二度計算されていますね。
nの値が大きいほどこの重複は増えていきます。

これを解決するアルゴリズムがあります。
DP（動的計画法）と言います。

fun fibonacci(n: Int): Long {
    if (n == 1) return 1L
    if (n == 2) return 1L

    val fibs = LongArray(n + 1)
    fibs[1] = 1L
    fibs[2] = 1L
    for (i in 3..n) {
        fibs[i] = fibs[i - 1] + fibs[i - 2]
    }

    return fibs[n]
}

fibsという配列に、n=1, n=2の値を初期値として詰めて、nに至るまで繰り返しi-1番目とi-2番目の要素を取り出して足して、再びfibsに値を詰めるという操作をしています。
無駄な計算がなく、とても綺麗だと思いました。
アルゴリズムを考える人たちってすごいですね。

最後に、弊社bitFlyerではAndroidエンジニアを募集しています。ご興味のある方はぜひご応募ください！