繰り返し二乗法と逆元をKotlinで実装

はじめに

KotlinでAtCoderに取り組んでいます。

Kotlin標準のMath.powはdouble or float型にて計算するため扱いにくいのでLong型対応のものを実装しました。(参考)

modを考慮しないpow

繰り返し二乗法により実装します。

fun pow(xx: Long, nn: Long): Long {
    var res = 1L
    var x = xx
    var n = nn
    while (n > 0) {
        if (n % 2 == 1L) res *= x
        x *= x
        n = n.shr(1)
    }
    return res
}

modを考慮したpow

掛け算のたびにmodをとります。

fun modPow(xx: Long, nn: Long, mod: Long): Long {
    var res = 1L
    var x = xx % mod
    var n = nn
    while (n > 0) {
        if (n % 2 == 1L) res = res * x % mod
        x = x * x % mod
        n = n.shr(1)
    }
    return res
}

フェルマーの小定理による逆元の算出

modPowを使用します。

fun modInv(x: Long, mod: Long): Long {
    return modPow(x, mod - 2, mod)
}

拡張Euclidの互除法による逆元の算出

fun modInv(xx: Long, mod: Long): Long {
    var x = xx
    var b = mod
    var u = 1L
    var v = 0L
    while (b > 0) {
        val t = x / b
        x -= t * b
        x = b.also {b = x}
        u -= t * v
        u = v.also {v = u}
    }
    u %= mod
    if (u < 0) u += mod
    return u
}