Rustのジェネリクスパラメーターに定数を渡せるようになる(const generics) - Qiita をアドベントカレンダーで読んだのですが、const genericsでできることってなんだろうなぁと思っていたら思い出したので記事にした次第です。

modint構造体

競プロでは素数$10^9 + 7$の余りを回答する問題があります。そんな時に使われるのがmodint構造体です。説明はリンク先が詳しいです。(ついでに二項係数のリンクも載せます)

C++の実装ではジェネリクスに整数を指定出来ていますが、今回Rust 1.50でstableになるconst genericsはまさにこの機能です！

というわけで、$10^9 + 7$以外の素数にしたいときなどに、C++同様にmodintの実装部分とは切り離して素数を指定できるようになったので、それを確認してみます。

ぶっちゃけ動的に指定できる(入力で素数が与えられるなど)ようになるわけでは相変わらずないので、コードが多少美しくなっただけかも()

modint/lib.rs

use std::fmt;
use std::ops;

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct ModInt<const MOD: usize> {
    val: usize,
}

impl<const MOD: usize> fmt::Display for ModInt<MOD> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.val)
    }
}

impl<const MOD: usize> ModInt<MOD> {
    pub fn new(n: usize) -> Self {
        Self { val: n % MOD }
    }

    pub fn val(&self) -> usize {
        // 念のためMOD演算
        self.val % MOD
    }

    pub fn _set_val(&mut self, val: usize) {
        self.val = val % MOD;
    }

    pub fn pow_u(&self, mut n: usize) -> Self {
        let mut val = self.val;
        let mut res: usize = 1;
        while n > 0 {
            if n % 2 == 1 {
                res = (res * val) % MOD;
            }
            val = (val * val) % MOD;
            n /= 2;
        }

        Self { val: res }
    }

    pub fn pow(&self, other: Self) -> Self {
        self.pow_u(other.val)
    }

    pub fn inv(&self) -> Self {
        self.pow_u(MOD - 2)
    }
}

impl<const MOD: usize> ops::Add for ModInt<MOD> {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        Self {
            val: (self.val + other.val) % MOD,
        }
    }
}

impl<const MOD: usize> ops::AddAssign for ModInt<MOD> {
    fn add_assign(&mut self, other: Self) {
        *self = Self {
            val: (self.val + other.val) % MOD,
        };
    }
}

impl<const MOD: usize> ops::Mul for ModInt<MOD> {
    type Output = Self;

    fn mul(self, other: Self) -> Self {
        Self {
            val: (self.val * other.val) % MOD,
        }
    }
}

impl<const MOD: usize> ops::MulAssign for ModInt<MOD> {
    fn mul_assign(&mut self, other: Self) {
        *self = Self {
            val: (self.val * other.val) % MOD,
        };
    }
}

impl<const MOD: usize> ops::Sub for ModInt<MOD> {
    type Output = Self;

    fn sub(mut self, other: Self) -> Self {
        if self.val < other.val {
            self.val += MOD;
        }
        Self {
            val: self.val - other.val % MOD,
        }
    }
}

impl<const MOD: usize> ops::SubAssign for ModInt<MOD> {
    fn sub_assign(&mut self, other: Self) {
        if self.val < other.val {
            self.val += MOD;
        }
        *self = Self {
            val: (self.val - other.val) % MOD,
        };
    }
}

impl<const MOD: usize> ops::Div for ModInt<MOD> {
    type Output = Self;

    fn div(self, other: Self) -> Self {
        if other.val == 0 {
            panic!("0 division occured.");
        }

        self * other.inv()
    }
}

impl<const MOD: usize> ops::DivAssign for ModInt<MOD> {
    fn div_assign(&mut self, other: Self) {
        if other.val == 0 {
            panic!("0 division occured.");
        }

        *self *= other.inv();
    }
}

pub struct ModCom<const MOD: usize> {
    fac: Vec<usize>,
    finv: Vec<usize>,
}

impl<const MOD: usize> ModCom<MOD> {
    pub fn new(cap: usize) -> Self {
        let mut fac = vec![0; cap];
        let mut finv = vec![0; cap];
        let mut inv = vec![0; cap];
        fac[0] = 1;
        fac[1] = 1;
        finv[0] = 1;
        finv[1] = 1;
        inv[1] = 1;
        for i in 2..cap {
            fac[i] = fac[i - 1] * i % MOD;
            inv[i] = MOD - inv[MOD % i] * (MOD / i) % MOD;
            finv[i] = finv[i - 1] * inv[i] % MOD;
        }

        Self { fac, finv }
    }

    pub fn com(&self, n: usize, k: usize) -> usize {
        if n < k {
            return 0;
        }
        self.fac[n] * (self.finv[k] * self.finv[n - k] % MOD) % MOD
    }
}

#[cfg(test)]
mod tests {
    use super::modint::*;
    type MINT = ModInt<1_000_000_007>;

    #[test]
    fn test1() {
        let a = MINT::new(111);
        let b = MINT::new(222);
        let c = MINT::new(333);
        let d = MINT::new(444);

        let res = a * b + c - d;
        assert_eq!(res.val(), 24531);
    }

    #[test]
    fn test2() {
        let a = MINT::new(111111111);
        let b = MINT::new(222222222);
        let c = MINT::new(333333333);
        let d = MINT::new(444444444);

        let res = a * b + c - d;
        assert_eq!(res.val(), 691358032);
    }
}

使用例

const genericsはRust 1.50以降の機能です。したがって、記事投稿時現在(2021/1/7)、const genericsはまだAtCoder(Rust 1.42.0)では使用できません。

...が、AtCoderの問題で使用例を出しておきます。当然ACかどうかなんて検証してません。

B - Training Camp

Rust

use modint::*;
use proconio::input;

const MOD: usize = 1_000_000_007;

fn main() {
    input! {
        n: usize,
    }

    let mut res = ModInt::<MOD>::new(1);

    for i in 1..=n {
        // 型推論される
        res *= ModInt::new(i);
    }

    println!("{}", res);
}

D - 多重ループ

Rust

use modint::*;
use proconio::input;

const MOD: usize = 1_000_000_007;

fn main() {
    input! {
        n: usize,
        k: usize,
    }

    let modcom = ModCom::<MOD>::new(n + k);
    println!("{}", modcom.com(n + k - 1, k));
}

const genericsがなかった時と比べれば、きれいに書けるようにはなったと思います。

実行方法

下記のように+nightlyフラグを立ててください。

bash

$ cargo +nightly --version
cargo 1.50.0-nightly (75d5d8cff 2020-12-22)
$ cargo +nightly run

あるいはrust-toolchainを使う方法もあるようです。

動的に(普通に)法を指定できるようにすると型システムの恩恵を受けにくくなり、今回みたいに静的だと利便性に難がありそうです。動的ディスパッチなんかで型として定義できつつ実行時指定できたら面白そうですが、あり得なさそう...?

modint以外でconst genericsのユースケースがちょっと思い浮かばなかったので、もし「こういう使い方があるよ！」みたいなものがあればコメント欄などで教えていただけると幸いです。(他の使い方は配列のサイズ( <const SIZE: usize> で [T; SIZE] )などでしょうか...?)

ここまで読んでいただきありがとうございました！

Rustでmodint構造体 by const generics

modint構造体

使用例

B - Training Camp

D - 多重ループ

実行方法