結果

問題 No.3350 Tree and Two Apples
コンテスト
ユーザー Solalyth
提出日時 2025-11-13 17:15:25
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 160 ms / 2,000 ms
コード長 8,707 bytes
コンパイル時間 13,927 ms
コンパイル使用メモリ 399,512 KB
実行使用メモリ 28,004 KB
最終ジャッジ日時 2025-11-13 21:29:36
合計ジャッジ時間 19,407 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 35
権限があれば一括ダウンロードができます

ソースコード

diff # 

#![allow(unused_must_use, non_snake_case, unused_labels, unused_imports, non_upper_case_globals)]
use external::*;
use cplib::prelude::*;
use nest as vec;
const INTERACTIVE: bool = false;
// const INTERACTIVE: bool = true; use input_interactive as input;

fn solve() {
    input! {
        N: usize, UV: [(usize1, usize1); N-1]
    }
    
    const M: u64 = 998244353;
    // 線形合同法による疑似乱数
    let mut lcg = std::iter::successors(Some((230183, 13182938)), |&(x, y)| {
        Some(((x*30253231+230183) % M, (y*23018237+2018239) % M))
    });
    let mut pool = HashMap::new();
    let mut f = |e| {
        *pool.entry(e).or_insert_with(|| {
            lcg.next().unwrap()
        })
    };
    
    
    
    let mut E = vec![void; N];
    for &(u, v) in &UV {
        E[u].push(v);
        E[v].push(u);
    }
    
    let mut par = vec![0; N];
    let mut order = vec![0];
    for i in 0..N {
        let i = order[i];
        for &j in &E[i] {
            if j == par[i] { continue; }
            par[j] = i;
            order.push(j);
        }
    }
    
    let mut up = vec![(0, 0); N];
    let mut up2 = vec![0; N];
    let mut set = vec![];
    for &i in order[1..].iter().rev() {
        for &j in &E[i] {
            if j == par[i] { continue; }
            up[i] = (up[i].0 + up[j].0, up[i].1 + up[j].1);
            set.push((up[j], up2[j]));
        }
        set.sort(); set.dedup();
        up[i] = f((up[i].0 % M, up[i].1 % M));
        up2[i] = set.iter().map(|e| e.1).sum::<u64>()+1;
        set.clear();
    }
    
    let mut down = vec![(0, 0); N];
    let mut down2 = vec![0; N];
    let mut ans = vec![];
    let mut map = HashMap::new();
    for &i in &order {
        let mut sum = (0, 0);
        let mut cnt = 1;
        for &j in &E[i] {
            if par[i] == j {
                sum = (sum.0 + down[i].0, sum.1 + down[i].1);
                let e = map.entry((down[i], down2[i])).or_insert(0);
                if *e == 0 { cnt += down2[i]; }
                *e += 1;
            } else {
                sum = (sum.0 + up[j].0, sum.1 + up[j].1);
                let e = map.entry((up[j], up2[j])).or_insert(0);
                if *e == 0 { cnt += up2[j]; }
                *e += 1;
            }
        }
        
        ans.push(((sum.0%M, sum.1%M), cnt));
        
        for &j in &E[i] {
            if par[i] == j { continue; }
            let (sum, mut cnt) = ((sum.0-up[j].0, sum.1-up[j].1), cnt);
            if map[&(up[j], up2[j])] == 1 {
                cnt -= up2[j];
            }
            down[j] = f((sum.0 % M, sum.1 % M));
            down2[j] = cnt;
        }
        map.clear();
    }
    
    ans.sort();
    ans.dedup();
    
    out << ans.iter().map(|e| e.1).sum::<u64>();
}




fn main() { out::init(INTERACTIVE || !SUBMISSION); solve(); out::print() }

mod external { pub use { proconio::{ input, input_interactive, marker::{Chars as chars, Usize1 as usize1} }, }; }

// my library: https://github.com/solalyth/cplib-rs
mod cplib { #![allow(unused_macros, dead_code)] pub const SUBMISSION: bool = true; pub mod prelude { pub use std::{ collections::{VecDeque, HashMap, HashSet, BTreeMap, BTreeSet, BinaryHeap}, cmp::{Ordering, Reverse}, mem::{swap, replace} }; pub use crate::cplib::{ *, SUBMISSION, util::{ output::{out, end} }, }; use std::fmt; #[derive(Clone, Copy, PartialEq, PartialOrd)] pub struct F64(pub f64); impl Eq for F64 {} impl Ord for F64 { fn cmp(&self, other: &Self) -> Ordering { self.partial_cmp(other).unwrap() } } impl fmt::Debug for F64 { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{}", self.0) } } } pub mod debug { pub fn replace_inf(s: &str) -> String { let (mut res, mut stk) = (String::new(), String::new()); for c in s.chars().chain(std::iter::once('*')) { if c.is_numeric() || c == '.' { stk.push(c); continue; } if stk.parse::<f64>().map_or(false, |s| 1.15e18 <= s) { res += "inf"; } else { res += &stk; } stk.clear(); res.push(c); } res.pop(); res } #[macro_export] macro_rules! epr { ($($args:tt)*) => { if !$crate::SUBMISSION { eprintln!("\x1b[31m{}\x1b[0m", crate::debug::replace_inf(&format!($($args)*))); } } } #[macro_export] macro_rules! table { ($vv:expr) => { if !$crate::SUBMISSION { use std::fmt::Write; let (mut lmax, mut res) = (0, String::new()); let vv: Vec<Vec<String>> = $vv.iter().map(|v| v.iter().map(|e| { let t = crate::debug::replace_inf(&format!("{e:?}")); lmax = lmax.max(t.len()); t }).collect()).collect(); for v in &vv { res += "    [ "; for e in v { write!(&mut res, "{e: >lmax$}, "); } res.pop(); res.pop(); res += "]\n"; } eprintln!("\x1b[38;5;208m{} = [\n{res}]\x1b[0m", stringify!($vv)); } }; } #[macro_export] macro_rules! oj_local { ($oj:expr, $local:expr) => { if $crate::SUBMISSION { $oj } else { $local } }; } } pub mod ds { } pub mod algo { } pub mod graph { } pub mod math { } pub mod mod998 { } pub mod util { pub mod output { #![allow(static_mut_refs, non_camel_case_types)] use std::{mem::replace, ops::{Not, Shl}, fmt::Write}; static mut BUFFER: Buffer = Buffer { buf: String::new(), endp: false, prev: Previous::LineHead }; pub struct Buffer { buf: String, endp: bool, prev: Previous, } impl Buffer { const LEN: usize = 16*1024*1024; fn print(&mut self) { if replace(&mut self.prev, Previous::LineHead) == Previous::LineHead { self.buf.pop(); } if crate::cplib::SUBMISSION { println!("{}", self.buf); } else { eprint!("\x1b[32m"); if self.buf.is_empty() { println!(">> (empty)"); } else { for s in self.buf.split('\n') { eprint!(">> "); println!("{s}"); } } eprint!("\x1b[0m"); } self.buf.clear(); } fn space(&mut self, sp: bool) { let prev = replace(&mut self.prev, if sp {Previous::Space} else {Previous::NoSpace}); if (sp || prev == Previous::Space) && prev != Previous::LineHead { self.buf.push(' '); } } } #[derive(Clone, Copy)] pub struct out; pub struct out_usp; pub struct end; #[derive(PartialEq, Eq, Clone, Copy)] enum Previous { Space, NoSpace, LineHead, } impl out { pub fn init(endp: bool) { unsafe { BUFFER.buf.reserve(Buffer::LEN); BUFFER.endp = endp; } } pub fn print() { unsafe { BUFFER.print(); } } pub fn space() { unsafe { if BUFFER.prev == Previous::NoSpace { BUFFER.prev = Previous::Space; } } } fn push<T: Primitive>(v: &T) { unsafe { BUFFER.space(true); v.fmt(&mut BUFFER.buf); } } } impl out_usp { fn push<T: Primitive>(v: &T) { unsafe { BUFFER.space(false); v.fmt(&mut BUFFER.buf); } } } impl Not for out { type Output = out_usp; fn not(self) -> Self::Output { out_usp } } macro_rules! impl_outs { ($($t:ty),+) => { $( impl<T: Primitive> Shl<T> for $t { type Output = Self; fn shl(self, rhs: T) -> Self::Output { Self::push(&rhs); self } } impl Shl<end> for $t { type Output = Self; fn shl(self, _: end) -> Self::Output { unsafe { if BUFFER.endp { BUFFER.print(); } else { BUFFER.buf += "\n"; BUFFER.prev = Previous::LineHead; } } self } } )+ }; } impl_outs!(out, out_usp); macro_rules! impl_for_slices { ($($t:ty),+) => { $(impl_for_slices!($t; out, out_usp);)+ }; ($t:ty; $($u:ty),+) => { $( impl<T: Primitive> Shl<$t> for $u { type Output = Self; fn shl(self, rhs: $t) -> Self::Output { for v in rhs { Self::push(v); } self } } )+} } impl_for_slices!(&[T], &Vec<T>); trait Primitive { fn fmt(&self, buf: &mut String); } macro_rules! impl_primitive { ($($t:ty),+) => { $( impl Primitive for $t { fn fmt(&self, buf: &mut String) { write!(buf, "{self}").ok(); } } )+ } } impl_primitive!(char, u32, u64, u128, usize, i32, i64, i128, f32, f64, &str, &String, String); impl Primitive for u8 { fn fmt(&self, buf: &mut String) { buf.push(*self as char); } } impl Primitive for bool { fn fmt(&self, buf: &mut String) { *buf += if *self { "Yes" } else { "No" }; } } } pub mod macros { #[macro_export] macro_rules! nest { [void; $n:expr] => { std::vec![std::vec![]; $n] }; [void; $n:expr $(;$m:expr)+] => { std::vec![crate::nest![void$(;$m)+]; $n] }; [$($v:expr),*] => { std::vec![$($v),*] }; [$e:expr; $n:expr] => { std::vec![$e; $n] }; [$e:expr; $n:expr $(;$m:expr)+] => { std::vec![crate::nest![$e$(;$m)+]; $n] }; } #[macro_export] macro_rules! min { ($($vl:expr),+) => { [$($vl),+].into_iter().reduce(|x,y| if x <= y {x} else {y}).unwrap() } } #[macro_export] macro_rules! max { ($($vl:expr),+) => { [$($vl),+].into_iter().reduce(|x,y| if x >= y {x} else {y}).unwrap() } } #[macro_export] macro_rules! chmin { ($dst:expr; $v:expr) => { { let v = $v; if v < $dst { $dst = v; true } else { false } } }; ($dst:expr; $($vl:expr),+) => { crate::chmin!($dst; crate::min!($($vl),+)) } } #[macro_export] macro_rules! chmax { ($dst:expr; $v:expr) => { { let v = $v; if $dst < v { $dst = v; true } else { false } } }; ($dst:expr; $($vl:expr),+) => { crate::chmax!($dst; crate::max!($($vl),+)) } } } } }
0