結果

問題 No.3250 最小公倍数
コンテスト
ユーザー LyricalMaestro
提出日時 2025-10-11 18:05:32
言語 Rust
(1.83.0 + proconio)
結果
TLE  
実行時間 -
コード長 4,550 bytes
コンパイル時間 14,292 ms
コンパイル使用メモリ 398,280 KB
実行使用メモリ 361,876 KB
最終ジャッジ日時 2025-10-16 17:10:45
合計ジャッジ時間 39,756 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 19 TLE * 3
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: unused imports: `Read` and `self`
 --> src/main.rs:3:15
  |
3 | use std::io::{self, Read};
  |               ^^^^  ^^^^
  |
  = note: `#[warn(unused_imports)]` on by default

ソースコード

diff # 

use proconio::input;
use std::collections::{HashMap, HashSet, VecDeque};
use std::io::{self, Read};

const MOD: u64 = 998_244_353;

fn mod_pow(mut base: u64, mut exp: u64, modu: u64) -> u64 {
    let mut result = 1u64;
    base %= modu;
    while exp > 0 {
        if exp % 2 == 1 {
            result = result * base % modu;
        }
        base = base * base % modu;
        exp /= 2;
    }
    result
}

fn main() {
    // 入力全読み取り
    input! {
        n: usize,
        a: [u64; n],
        uv: [(usize, usize); n -1]
    };

    // 隣接リスト
    let mut next_nodes = vec![Vec::new(); n];
    for j in 0..n - 1 {
        let (u, v) = uv[j];
        next_nodes[u - 1].push(v - 1);
        next_nodes[v - 1].push(u - 1);
    }

    let sqrt_max_a = (a.iter().copied().max().unwrap() as f64).sqrt() as usize;

    // 素数リスト生成（単純な篩）
    let mut primes = Vec::new();
    let mut is_prime = vec![true; sqrt_max_a + 1];
    for p in 2..=sqrt_max_a {
        if !is_prime[p] {
            continue;
        }
        primes.push(p as u64);
        let mut x = 2 * p;
        while x <= sqrt_max_a {
            is_prime[x] = false;
            x += p;
        }
    }

    // 各ノードごとの因数分解
    let mut a1 = Vec::with_capacity(n);
    let mut a2 = Vec::with_capacity(n);
    for mut val in a.clone() {
        let mut map = HashMap::new();
        for &p in &primes {
            while val % p == 0 {
                *map.entry(p).or_insert(0u64) += 1;
                val /= p;
            }
        }
        a1.push(val);
        a2.push(map);
    }

    // DFS（スタックによる非再帰）
    let mut stack = VecDeque::new();
    let mut parents = vec![-2; n];
    parents[0] = -1;
    stack.push_back((0usize, 0usize));

    let mut answer1 = vec![1u64; n];
    let mut answer2: Vec<HashMap<u64, u64>> = vec![HashMap::new(); n];
    let mut answer1_set: Vec<HashSet<u64>> = vec![HashSet::new(); n];

    while let Some((v, mut index)) = stack.pop_back() {
        if index == 0 {
            // 初回訪問時
            answer1[v] = (answer1[v] * a1[v]) % MOD;
            answer1_set[v].insert(a1[v]);

            for (&p, &v_ind) in &a2[v] {
                answer2[v].entry(p).and_modify(|x| *x = (*x).max(v_ind)).or_insert(v_ind);
            }
        }

        while index < next_nodes[v].len() {
            let w = next_nodes[v][index];
            if w as i64 == parents[v] {
                index += 1;
                continue;
            }
            parents[w] = v as i64;
            stack.push_back((v, index + 1));
            stack.push_back((w, 0));
            break;
        }

        if index == next_nodes[v].len() {
            let p = parents[v];
            if p != -1 {
                let p = p as usize;

                // answer1_set と answer1 の統合
                let mut a_set1 = std::mem::take(&mut answer1_set[p]);
                let mut a_set2 = std::mem::take(&mut answer1_set[v]);
                if a_set1.len() >= a_set2.len() {
                    let mut value = answer1[p];
                    for &prime in &a_set2 {
                        if !a_set1.contains(&prime) {
                            value = (value * prime) % MOD;
                            a_set1.insert(prime);
                        }
                    }
                    answer1[p] = value;
                    answer1_set[p] = a_set1;
                } else {
                    let mut value = answer1[v];
                    for &prime in &a_set1 {
                        if !a_set2.contains(&prime) {
                            value = (value * prime) % MOD;
                            a_set2.insert(prime);
                        }
                    }
                    answer1[p] = value;
                    answer1_set[p] = a_set2;
                }

                // answer2 の統合
                let v_map = answer2.get(v).unwrap().iter().map(|(p_, v_ind)| (*p_, *v_ind)).collect::<Vec<_>>();
                let a_map = answer2.get_mut(p).unwrap();
                for (p_, v_ind) in v_map.iter() {
                    a_map
                        .entry(*p_)
                        .and_modify(|x| *x = (*x).max(*v_ind))
                        .or_insert(*v_ind);
                }
            }
        }
    }

    for i in 0..n {
        let mut ans = answer1[i];
        for (&p, &v_ind) in &answer2[i] {
            ans = ans * mod_pow(p, v_ind, MOD) % MOD;
        }
        println!("{}", ans);
    }
}
0