Submission #55678131

---

Source Code Expand

use std::collections::VecDeque;

use itertools::Itertools;
use proconio::{input, marker::Usize1};

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

    let edge = {
        let mut edge = vec![Vec::new(); n];
        for _ in 1..n {
            input! { x: Usize1, y: Usize1 }
            edge[x].push(y);
            edge[y].push(x);
        }
        edge
    };

    // DFS starting at the node 0
    let mut dp = {
        let mut dp = vec![0; n];
        let mut stack = Vec::with_capacity(n);
        stack.push(0);
        let mut is_visited_once = vec![false; n];
        while let Some(&i) = stack.last() {
            if is_visited_once[i] {
                // postorder
                is_visited_once[i] = false;
                stack.pop();

                dp[i] = edge[i]
                    .iter()
                    .filter(|&&j| !is_visited_once[j])
                    .fold(1, |acc, &j| acc * (dp[j] + 1) % m);
            } else {
                // preorder
                is_visited_once[i] = true;
                for &j in edge[i].iter().filter(|&&j| !is_visited_once[j]) {
                    stack.push(j)
                }
            }
        }
        dp
    };
    // println!("DFS starting at 0: {:?}\n", dp);

    // Rerooting: BFS starting at the node 0
    let mut deque = VecDeque::with_capacity(n);
    deque.push_back((0, 1));
    let mut is_visited = vec![false; n];
    while let Some((i, left)) = deque.pop_front() {
        is_visited[i] = true;

        let children = Vec::from_iter(edge[i].iter().filter(|&&j| !is_visited[j]).copied());
        // println!("children: {:?}", children);

        let mut acc_l = vec![1; children.len()];
        for i in 1..children.len() {
            acc_l[i] = acc_l[i - 1] * (dp[children[i - 1]] + 1) % m
        }
        let mut acc_r = vec![1; children.len()];
        for i in (1..children.len()).rev() {
            acc_r[i - 1] = acc_r[i] * (dp[children[i]] + 1) % m
        }
        // println!("{:?}\n{:?}", acc_l, acc_r);

        for (i, c) in children.into_iter().enumerate() {
            let left = acc_l[i] * acc_r[i] % m * left % m + 1;
            dp[c] = (dp[c] * left) % m;
            deque.push_back((c, left))
        }
        // println!("{:?}\n", dp)
    }

    println!("{}", dp.into_iter().join("\n"))
}

Submission Info

Judge Result