AtCoder Beginner Contest 300 E

https://atcoder.jp/contests/abc300/tasks/abc300_e

例えば6なら、確率をで表すとすると、

なので、

となります。一般には、

と書けます。

メモを関数の外に持つのはクレートを使わないと難しそう。

// Dice Product 3
#![allow(non_snake_case)]

use std::collections::HashMap;


//////////////////// constants ////////////////////

const D: u64 = 998244353;


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

// ax = by + 1 (a, b > 0)
fn linear_diophantine(a: i64, b: i64) -> Option<(i64, i64)> {
    if a == 1 {
        return Some((1, 0))
    }
    
    let q = b / a;
    let r = b % a;
    if r == 0 {
        return None
    }
    let (x1, y1) = linear_diophantine(r, a)?;
    Some((-q * x1 - y1, -x1))
}

fn inverse(a: i64, d: i64) -> i64 {
    let (x, _y) = linear_diophantine(a, d).unwrap();
    if x >= 0 {
        x % d
    }
    else {
        x % d + d
    }
}


//////////////////// process ////////////////////

fn f(N: u64, memo: &mut HashMap<u64, u64>) -> u64 {
    if N == 1 {
        1
    }
    else if memo.get(&N) != None {
        *memo.get(&N).unwrap()
    }
    else {
        let inv5 = inverse(5, D as i64) as u64;
        let x = (2..7).filter(|d| N % d == 0).
                    map(|d| f(N/d, memo)).sum::<u64>() * inv5 % D;
        memo.insert(N, x);
        x
    }
}

//////////////////// main ////////////////////

fn main() {
    let N: u64 = read();
    let mut memo = HashMap::new();
    println!("{}", f(N, &mut memo))
}