#1116347 (Rust) No.3250 最小公倍数

• 問題: No.3250 最小公倍数
• ユーザー: urectanc
• 提出日時: 2025-08-30 16:07:05
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 1,319 ms / 2,000 ms
• コード長: 16,703 bytes
• コンパイル時間: 11,780 ms
• コンパイル使用メモリ: 400,844 KB
• 実行使用メモリ: 139,776 KB
• 最終ジャッジ日時: 2025-08-30 16:07:37
• 合計ジャッジ時間: 24,257 ms
• ジャッジサーバーID （参考情報）: judge5 / judge2

ファイルパターン	結果
sample	AC * 1
other	AC * 21

ソースコード

use std::io::Write;

use modint::ModInt998244353;
use proconio::{input, marker::Usize1};

type Mint = ModInt998244353;

fn main() {
    input! {
        n: usize,
        a: [usize; n],
        edges: [(Usize1, Usize1); n - 1],
    }
    let mut stdout = std::io::BufWriter::new(std::io::stdout().lock());
    let a_max = a.iter().max().copied().unwrap();

    let mut lpf = (0..=a_max).collect::<Vec<_>>();
    for i in 2..=a_max {
        if lpf[i] == i {
            for j in (i * i..=a_max).step_by(i) {
                if lpf[j] == j {
                    lpf[j] = i;
                }
            }
        }
    }

    let factorize = |mut x: usize| {
        let mut res = vec![];

        while x > 1 {
            let p = lpf[x];
            let mut cnt = 0;
            while x % p == 0 {
                x /= p;
                cnt += 1;
            }
            res.push((p, cnt));
        }

        res
    };

    let factorized = a.iter().map(|&a| factorize(a)).collect::<Vec<_>>();

    let mut ans = vec![Mint::new(0); n];
    let mut lcm = Lcm::new(a_max);

    let mut graph = vec![vec![]; n];
    for &(u, v) in &edges {
        graph[u].push(v);
        graph[v].push(u);
    }

    let mut pre_order = vec![];
    let mut stack = vec![(0, !0)];
    while let Some((v, p)) = stack.pop() {
        pre_order.push(v);
        graph[v].retain(|&u| u != p);
        for &u in &graph[v] {
            stack.push((u, v));
        }
    }

    let mut subtree_size = vec![0; n];
    for &v in pre_order.iter().rev() {
        if !graph[v].is_empty() {
            graph[v].select_nth_unstable_by_key(0, |&u| std::cmp::Reverse(subtree_size[u]));
        }
        subtree_size[v] = graph[v].iter().map(|&u| subtree_size[u]).sum::<usize>() + 1;
    }

    fn dfs(
        graph: &Vec<Vec<usize>>,
        subtree_size: &Vec<usize>,
        ans: &mut Vec<Mint>,
        factorized: &Vec<Vec<(usize, u64)>>,
        lcm: &mut Lcm,
        v: usize,
        keep: bool,
    ) {
        for &u in graph[v].iter().skip(1) {
            dfs(graph, subtree_size, ans, factorized, lcm, u, false);
        }
        if let Some(&u) = graph[v].get(0) {
            dfs(graph, subtree_size, ans, factorized, lcm, u, true);
        }

        for &u in graph[v].iter().skip(1) {
            dfs2(graph, subtree_size, ans, factorized, lcm, u);
        }
        lcm.add(&factorized[v]);

        ans[v] = lcm.val();

        if !keep {
            lcm.reset();
        }
    }

    fn dfs2(
        graph: &Vec<Vec<usize>>,
        subtree_size: &Vec<usize>,
        ans: &mut Vec<Mint>,
        factorized: &Vec<Vec<(usize, u64)>>,
        lcm: &mut Lcm,
        v: usize,
    ) {
        lcm.add(&factorized[v]);
        for &u in &graph[v] {
            dfs2(graph, subtree_size, ans, factorized, lcm, u);
        }
    }

    dfs(
        &graph,
        &subtree_size,
        &mut ans,
        &factorized,
        &mut lcm,
        0,
        false,
    );

    for ans in &ans {
        writeln!(stdout, "{ans}").unwrap();
    }
}

struct Lcm {
    pf_count: Vec<u64>,
    modified: Vec<usize>,
}

impl Lcm {
    fn new(max: usize) -> Self {
        Self {
            pf_count: vec![0; max + 1],
            modified: vec![],
        }
    }

    fn add(&mut self, factors: &[(usize, u64)]) {
        for &(p, cnt) in factors {
            if self.pf_count[p] == 0 {
                self.modified.push(p);
            }
            self.pf_count[p] = self.pf_count[p].max(cnt);
        }
    }

    fn reset(&mut self) {
        for p in self.modified.drain(..) {
            self.pf_count[p] = 0;
        }
    }

    fn val(&self) -> Mint {
        self.modified
            .iter()
            .map(|&p| Mint::new(p as u64).pow(self.pf_count[p]))
            .product()
    }
}

#[allow(dead_code)]
mod hld {
    pub struct HLD {
        parent: Vec<usize>,
        index: Vec<usize>,
        head: Vec<usize>,
        sorted: Vec<usize>,
        subtree_size: Vec<usize>,
    }

    impl HLD {
        pub fn from_edges(
            root: usize,
            edges: impl ExactSizeIterator<Item = (usize, usize)>,
        ) -> Self {
            let n = edges.len() + 1;
            let mut graph = vec![vec![]; n];
            for (u, v) in edges {
                graph[u].push(v);
                graph[v].push(u);
            }

            let mut sorted = Vec::with_capacity(n);
            let mut parent = vec![!0; n];
            let mut stack = vec![root];
            while let Some(v) = stack.pop() {
                sorted.push(v);
                graph[v].retain(|&u| u != parent[v]);
                for &u in &graph[v] {
                    parent[u] = v;
                    stack.push(u);
                }
            }

            let mut subtree_size = vec![1; n];
            for &v in sorted.iter().rev() {
                if !graph[v].is_empty() {
                    graph[v].select_nth_unstable_by_key(0, |&c| std::cmp::Reverse(subtree_size[c]));
                }
                if v != root {
                    subtree_size[parent[v]] += subtree_size[v];
                }
            }

            let mut index = vec![!0; n];
            let mut head = (0..n).collect::<Vec<_>>();
            sorted.clear();
            stack.push(root);
            while let Some(v) = stack.pop() {
                index[v] = sorted.len();
                sorted.push(v);
                if let Some(&c) = graph[v].first() {
                    head[c] = head[v];
                }
                stack.extend(graph[v].iter().copied().rev());
            }

            Self {
                parent,
                index,
                head,
                sorted,
                subtree_size,
            }
        }

        pub fn sorted(&self) -> &'_ Vec<usize> {
            &self.sorted
        }

        pub fn parent(&self, v: usize) -> Option<usize> {
            Some(self.parent[v]).filter(|&v| v != !0)
        }

        pub fn index(&self, v: usize) -> usize {
            self.index[v]
        }

        pub fn edge_index(&self, u: usize, v: usize) -> usize {
            self.index[u].max(self.index[v])
        }

        pub fn subtree(&self, v: usize) -> (usize, usize) {
            let l = self.index[v];
            (l, l + self.subtree_size[v])
        }

        pub fn is_ancestor(&self, u: usize, v: usize) -> bool {
            let (l, r) = self.subtree(u);
            (l..r).contains(&self.index(v))
        }

        pub fn la(&self, mut v: usize, mut d: usize) -> usize {
            while v != !0 {
                let u = self.head[v];
                if self.index[v] - self.index[u] >= d {
                    v = self.sorted[self.index[v] - d];
                    break;
                }
                d -= self.index[v] - self.index[u] + 1;
                v = self.parent[u];
            }
            v
        }

        pub fn lca(&self, mut u: usize, mut v: usize) -> usize {
            if self.index(u) > self.index(v) {
                std::mem::swap(&mut u, &mut v);
            }

            if self.is_ancestor(u, v) {
                return u;
            }

            while self.index(u) < self.index(v) {
                v = self.parent[self.head[v]];
            }

            v
        }

        pub fn dist(&self, u: usize, v: usize) -> usize {
            self.path(u, v)
                .map(|(l, r, last)| usize::from(!last) + self.index[r] - self.index[l])
                .sum()
        }

        pub fn path_vertices(
            &self,
            u: usize,
            v: usize,
        ) -> impl Iterator<Item = (usize, usize)> + '_ {
            self.path(u, v)
                .map(|(u, v, _)| (self.index[u], self.index[v] + 1))
        }

        pub fn path_edges(&self, u: usize, v: usize) -> impl Iterator<Item = (usize, usize)> + '_ {
            self.path(u, v)
                .map(|(u, v, last)| (self.index[u] + usize::from(last), self.index[v] + 1))
        }

        fn path(&self, u: usize, v: usize) -> PathSegments<'_> {
            PathSegments {
                hld: self,
                u,
                v,
                exhausuted: false,
            }
        }
    }

    pub struct PathSegments<'a> {
        hld: &'a HLD,
        u: usize,
        v: usize,
        exhausuted: bool,
    }

    impl Iterator for PathSegments<'_> {
        type Item = (usize, usize, bool);

        fn next(&mut self) -> Option<Self::Item> {
            if self.exhausuted {
                return None;
            }

            let Self {
                hld:
                    HLD {
                        parent,
                        index,
                        head,
                        ..
                    },
                u,
                v,
                ..
            } = *self;

            if head[u] == head[v] {
                self.exhausuted = true;
                if index[u] < index[v] {
                    Some((u, v, true))
                } else {
                    Some((v, u, true))
                }
            } else {
                if index[u] < index[v] {
                    self.v = parent[head[v]];
                    Some((head[v], v, false))
                } else {
                    self.u = parent[head[u]];
                    Some((head[u], u, false))
                }
            }
        }
    }
}

#[allow(dead_code)]
mod modint {
    use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};

    pub type ModInt998244353 = ModInt<998244353>;
    pub type ModInt1000000007 = ModInt<1000000007>;

    type ModIntInner = u64;

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

    impl<const M: ModIntInner> ModInt<M> {
        const IS_PRIME: bool = is_prime(M as u32);

        pub const fn modulus() -> ModIntInner {
            M
        }

        pub const fn new(val: ModIntInner) -> Self {
            assert!(M < (1 << 31));
            Self {
                val: val.rem_euclid(M),
            }
        }

        pub const fn new_unchecked(val: ModIntInner) -> Self {
            Self { val }
        }

        pub const fn val(&self) -> ModIntInner {
            self.val
        }

        pub fn pow(self, mut exp: u64) -> Self {
            let mut result = Self::new(1);
            let mut base = self;
            while exp > 0 {
                if exp & 1 == 1 {
                    result *= base;
                }
                base *= base;
                exp >>= 1;
            }
            result
        }

        pub fn inv(self) -> Self {
            assert!(Self::IS_PRIME);
            self.pow(M as u64 - 2).into()
        }
    }

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

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

    impl<const M: ModIntInner> std::str::FromStr for ModInt<M> {
        type Err = std::num::ParseIntError;
        fn from_str(s: &str) -> Result<Self, Self::Err> {
            let value = s.parse::<ModIntInner>()?;
            Ok(ModInt::new(value))
        }
    }

    impl<const M: ModIntInner> Neg for ModInt<M> {
        type Output = Self;
        fn neg(mut self) -> Self::Output {
            if self.val > 0 {
                self.val = M - self.val;
            }
            self
        }
    }

    impl<const M: ModIntInner, T: Into<ModInt<M>>> AddAssign<T> for ModInt<M> {
        fn add_assign(&mut self, rhs: T) {
            self.val += rhs.into().val;
            if self.val >= M {
                self.val -= M;
            }
        }
    }

    impl<const M: ModIntInner, T: Into<ModInt<M>>> SubAssign<T> for ModInt<M> {
        fn sub_assign(&mut self, rhs: T) {
            self.val = self.val.wrapping_sub(rhs.into().val);
            if self.val > M {
                self.val = self.val.wrapping_add(M);
            }
        }
    }

    impl<const M: ModIntInner, T: Into<ModInt<M>>> MulAssign<T> for ModInt<M> {
        fn mul_assign(&mut self, rhs: T) {
            self.val = self.val * rhs.into().val % M;
        }
    }

    impl<const M: ModIntInner, T: Into<ModInt<M>>> DivAssign<T> for ModInt<M> {
        fn div_assign(&mut self, rhs: T) {
            *self *= rhs.into().inv();
        }
    }

    macro_rules! impl_binnary_operators {
    ($({ $op: ident, $op_assign: ident, $fn: ident, $fn_assign: ident}),*) => {$(
        impl<const M: ModIntInner, T: Into<ModInt<M>>> $op<T> for ModInt<M> {
            type Output = ModInt<M>;
            fn $fn(mut self, rhs: T) -> ModInt<M> {
                self.$fn_assign(rhs.into());
                self
            }
        }

        impl<const M: ModIntInner> $op<&ModInt<M>> for ModInt<M> {
            type Output = ModInt<M>;
            fn $fn(self, rhs: &ModInt<M>) -> ModInt<M> {
                self.$fn(*rhs)
            }
        }

        impl<const M: ModIntInner, T: Into<ModInt<M>>> $op<T> for &ModInt<M> {
            type Output = ModInt<M>;
            fn $fn(self, rhs: T) -> ModInt<M> {
                (*self).$fn(rhs.into())
            }
        }

        impl<const M: ModIntInner> $op<&ModInt<M>> for &ModInt<M> {
            type Output = ModInt<M>;
            fn $fn(self, rhs: &ModInt<M>) -> ModInt<M> {
                (*self).$fn(*rhs)
            }
        }

        impl<const M: ModIntInner> $op_assign<&ModInt<M>> for ModInt<M> {
            fn $fn_assign(&mut self, rhs: &ModInt<M>) {
                *self = self.$fn(*rhs);
            }
        }
    )*};
}

    impl_binnary_operators!(
        {Add, AddAssign, add, add_assign},
        {Sub, SubAssign, sub, sub_assign},
        {Mul, MulAssign, mul, mul_assign},
        {Div, DivAssign, div, div_assign}
    );

    impl<const M: ModIntInner> std::iter::Sum for ModInt<M> {
        fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
            iter.fold(Self::new(0), Add::add)
        }
    }

    impl<'a, const M: ModIntInner> std::iter::Sum<&'a Self> for ModInt<M> {
        fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
            iter.fold(Self::new(0), Add::add)
        }
    }

    impl<const M: ModIntInner> std::iter::Product for ModInt<M> {
        fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
            iter.fold(Self::new(1), Mul::mul)
        }
    }

    impl<'a, const M: ModIntInner> std::iter::Product<&'a Self> for ModInt<M> {
        fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
            iter.fold(Self::new(1), Mul::mul)
        }
    }

    macro_rules! impl_rem_euclid_signed {
    ($($ty:tt),*) => {
        $(
            impl<const M: ModIntInner> From<$ty> for ModInt<M> {
                fn from(value: $ty) -> ModInt<M> {
                    Self::new_unchecked((value as i64).rem_euclid(M as i64) as ModIntInner)
                }
            }
        )*
    };
}
    impl_rem_euclid_signed!(i8, i16, i32, i64, isize);

    macro_rules! impl_rem_euclid_unsigned {
    ($($ty:tt),*) => {
        $(
            impl<const M: ModIntInner> From<$ty> for ModInt<M> {
                fn from(value: $ty) -> ModInt<M> {
                    Self::new_unchecked((value as u64).rem_euclid(M as u64) as ModIntInner)
                }
            }
        )*
    };
}
    impl_rem_euclid_unsigned!(u8, u16, u32, u64, usize);

    const fn is_prime(n: u32) -> bool {
        const fn miller_rabin(n: u32, witness: u32) -> bool {
            let (n, witness) = (n as u64, witness as u64);
            let mut d = n >> (n - 1).trailing_zeros();
            let mut y = {
                let (mut res, mut pow, mut e) = (1, witness, d);
                while e > 0 {
                    if e & 1 == 1 {
                        res = res * pow % n;
                    }
                    pow = pow * pow % n;
                    e >>= 1;
                }
                res
            };
            while d != n - 1 && y != 1 && y != n - 1 {
                y = y * y % n;
                d <<= 1;
            }
            y == n - 1 || d & 1 == 1
        }

        const WITNESS: [u32; 3] = [2, 7, 61];

        if n == 1 || n % 2 == 0 {
            return n == 2;
        }

        let mut i = 0;
        while i < WITNESS.len() {
            if !miller_rabin(n, WITNESS[i]) {
                return false;
            }
            i += 1;
        }

        true
    }
}