#1161689 (Rust) No.2313 Product of Subsequence (hard)

• 問題: No.2313 Product of Subsequence (hard)
• ユーザー: LyricalMaestro
• 提出日時: 2026-04-26 03:26:31
• 言語: Rust (1.94.0 + proconio + num + itertools)
• 結果: TLE
• 実行時間: -
• コード長: 5,064 bytes
• コンパイル時間: 1,944 ms
• コンパイル使用メモリ: 206,308 KB
• 実行使用メモリ: 13,184 KB
• 最終ジャッジ日時: 2026-04-26 03:27:10
• 合計ジャッジ時間: 10,976 ms
• ジャッジサーバーID （参考情報）: judge1_1 / judge3_1

ファイルパターン	結果
sample	AC * 3
other	AC * 10 TLE * 1 -- * 16

ソースコード

use std::collections::HashMap;
use proconio::input;

const MOD: i64 = 998244353;

// ------------------------------
// Combination Calculator
// ------------------------------
struct CombinationCalculator {
    fact: Vec<i64>,
    inv_fact: Vec<i64>,
    modv: i64,
}

impl CombinationCalculator {
    fn new(size: usize, modv: i64) -> Self {
        let mut fact = vec![0; size + 1];
        fact[0] = 1;
        for i in 1..=size {
            fact[i] = fact[i - 1] * (i as i64) % modv;
        }

        let mut inv_fact = vec![0; size + 1];
        inv_fact[size] = mod_pow(fact[size], modv - 2, modv);

        for i in (0..size).rev() {
            inv_fact[i] = inv_fact[i + 1] * ((i + 1) as i64) % modv;
        }

        Self { fact, inv_fact, modv }
    }

    fn comb(&self, n: usize, r: usize) -> i64 {
        if n < r {
            return 0;
        }
        self.fact[n] * self.inv_fact[r] % self.modv * self.inv_fact[n - r] % self.modv
    }
}

// ------------------------------
// mod pow
// ------------------------------
fn mod_pow(mut base: i64, mut exp: i64, m: i64) -> i64 {
    let mut res = 1;
    base %= m;
    while exp > 0 {
        if exp & 1 == 1 {
            res = res * base % m;
        }
        base = base * base % m;
        exp >>= 1;
    }
    res
}

// ------------------------------
// DFSで状態列挙
// ------------------------------
fn dfs(
    k_primes_array: &Vec<(i64, usize)>,
    index: usize,
    p: i64,
    p_list: &mut Vec<usize>,
    p_to_list_map: &mut HashMap<i64, Vec<usize>>,
) {
    if index == k_primes_array.len() {
        p_to_list_map.insert(p, p_list.clone());
        return;
    }

    let (p0, value) = k_primes_array[index];
    let mut q = 1;

    for v in 0..=value {
        p_list.push(v);
        dfs(k_primes_array, index + 1, p * q, p_list, p_to_list_map);
        p_list.pop();
        q *= p0;
    }
}

fn main() {
    input! {
        n: usize,
        mut k: i64,
        a_vec: [i64; n],
    }

    if k == 1 {
        println!("{}", (mod_pow(2, n as i64, MOD) - 1 + MOD) % MOD);
        return;
    }

    let base_k = k;

    // ------------------------------
    // 素因数分解
    // ------------------------------
    let mut k_primes: HashMap<i64, usize> = HashMap::new();
    let mut p = 2;
    while p * p <= k {
        if k % p == 0 {
            let mut cnt = 0;
            while k % p == 0 {
                k /= p;
                cnt += 1;
            }
            k_primes.insert(p, cnt);
        }
        p += 1;
    }
    if k > 1 {
        k_primes.insert(k, 1);
    }

    let mut k_primes_array: Vec<(i64, usize)> =
        k_primes.into_iter().collect();
    k_primes_array.sort_by_key(|x| x.0);

    // ------------------------------
    // 状態列挙
    // ------------------------------
    let mut p_to_list_map: HashMap<i64, Vec<usize>> = HashMap::new();
    dfs(
        &k_primes_array,
        0,
        1,
        &mut Vec::new(),
        &mut p_to_list_map,
    );

    // ------------------------------
    // 遷移テーブル
    // ------------------------------
    let mut transition: HashMap<(i64, i64), i64> = HashMap::new();

    for (&from_p, from_list) in p_to_list_map.iter() {
        for (&trans_p, trans_list) in p_to_list_map.iter() {
            let mut q = 1;
            for j in 0..k_primes_array.len() {
                let (p, max_v) = k_primes_array[j];
                let exp = std::cmp::min(max_v, from_list[j] + trans_list[j]);
                q *= p.pow(exp as u32);
            }
            transition.insert((from_p, trans_p), q);
        }
    }

    // ------------------------------
    // Aを圧縮（freq）
    // ------------------------------
    let mut a_array: HashMap<i64, usize> = HashMap::new();

    for mut a in a_vec {
        let mut q = 1;
        for &(p, value) in &k_primes_array {
            let mut x = 0;
            while a % p == 0 {
                a /= p;
                x += 1;
            }
            x = std::cmp::min(value, x);
            q *= p.pow(x as u32);
        }
        *a_array.entry(q).or_insert(0) += 1;
    }

    let combi = CombinationCalculator::new(n, MOD);

    // ------------------------------
    // DP
    // ------------------------------
    let mut dp: HashMap<i64, i64> = HashMap::new();
    dp.insert(1, 1);

    for (&a, &freq_q) in a_array.iter() {
        let mut new_dp = dp.clone();

        for (&key_p, &value) in dp.iter() {
            let mut key_p0 = key_p;

            for q in 1..=freq_q {
                let new_key = transition[&(key_p0, a)];

                let add = value * combi.comb(freq_q, q) % MOD;

                *new_dp.entry(new_key).or_insert(0) += add;
                *new_dp.get_mut(&new_key).unwrap() %= MOD;

                key_p0 = new_key;
            }
        }

        dp = new_dp;
    }

    // ------------------------------
    // 出力
    // ------------------------------
    if let Some(&ans) = dp.get(&base_k) {
        println!("{}", ans);
    } else {
        println!("0");
    }
}