#669125 (Rust) No.1614 Majority Painting on Tree

提出ソース

結果

問題	No.1614 Majority Painting on Tree
ユーザー	ngtkana
提出日時	2021-06-23 02:11:25
言語	Rust (1.77.0 + proconio)
結果	AC
実行時間	1,130 ms / 5,000 ms
コード長	19,968 bytes
コンパイル時間	14,012 ms
コンパイル使用メモリ	378,596 KB
実行使用メモリ	5,376 KB
最終ジャッジ日時	2024-07-17 15:40:29
合計ジャッジ時間	26,844 ms
ジャッジサーバーID （参考情報）	judge1 / judge2

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	277 ms 5,248 KB
testcase_01	AC	101 ms 5,248 KB
testcase_02	AC	134 ms 5,376 KB
testcase_03	AC	220 ms 5,376 KB
testcase_04	AC	314 ms 5,376 KB
testcase_05	AC	299 ms 5,376 KB
testcase_06	AC	633 ms 5,376 KB
testcase_07	AC	446 ms 5,376 KB
testcase_08	AC	235 ms 5,376 KB
testcase_09	AC	611 ms 5,376 KB
testcase_10	AC	177 ms 5,376 KB
testcase_11	AC	730 ms 5,376 KB
testcase_12	AC	805 ms 5,376 KB
testcase_13	AC	575 ms 5,376 KB
testcase_14	AC	1,130 ms 5,376 KB
testcase_15	AC	1,081 ms 5,376 KB
testcase_16	AC	274 ms 5,376 KB
testcase_17	AC	522 ms 5,376 KB
testcase_18	AC	153 ms 5,376 KB
testcase_19	AC	75 ms 5,376 KB
testcase_20	AC	325 ms 5,376 KB
testcase_21	AC	333 ms 5,376 KB
testcase_22	AC	359 ms 5,376 KB
testcase_23	AC	57 ms 5,376 KB
testcase_24	AC	331 ms 5,376 KB
testcase_25	AC	93 ms 5,376 KB
testcase_26	AC	209 ms 5,376 KB
testcase_27	AC	1 ms 5,376 KB
testcase_28	AC	1 ms 5,376 KB
testcase_29	AC	1 ms 5,376 KB
testcase_30	AC	1 ms 5,376 KB
testcase_31	AC	1 ms 5,376 KB
testcase_32	AC	1 ms 5,376 KB
testcase_33	AC	1 ms 5,376 KB
testcase_34	AC	1 ms 5,376 KB
testcase_35	AC	1 ms 5,376 KB
testcase_36	AC	1 ms 5,376 KB
testcase_37	AC	1 ms 5,376 KB
testcase_38	AC	1 ms 5,376 KB
testcase_39	AC	1 ms 5,376 KB
testcase_40	AC	1 ms 5,376 KB
testcase_41	AC	1 ms 5,376 KB
testcase_42	AC	1 ms 5,376 KB
testcase_43	AC	1 ms 5,376 KB
testcase_44	AC	1 ms 5,376 KB
testcase_45	AC	1 ms 5,376 KB
testcase_46	AC	1 ms 5,376 KB
testcase_47	AC	1 ms 5,376 KB
testcase_48	AC	1 ms 5,376 KB

権限があれば一括ダウンロードができます

コンパイルメッセージ

warning: unused imports: `Leaf`, `Tuple`, `VecLen`
   --> src/main.rs:364:27
    |
364 |             multi_token::{Leaf, Parser, ParserTuple, RawTuple, Tuple, VecLen},
    |                           ^^^^                                 ^^^^^  ^^^^^^
    |
    = note: `#[warn(unused_imports)]` on by default

warning: unused import: `with_str`
   --> src/main.rs:620:35
    |
620 |     pub use self::i::{with_stdin, with_str};
    |                                   ^^^^^^^^

warning: unused imports: `ParserTuple`, `Parser`, `RawTuple`, `Token`, `Usize1`
   --> src/main.rs:623:28
    |
623 |         pub use super::i::{Parser, ParserTuple, RawTuple, Token, Usize1};
    |                            ^^^^^^  ^^^^^^^^^^^  ^^^^^^^^  ^^^^^  ^^^^^^

ソースコード

raw source code

use fp::F998244353 as Fp;
use std::iter::{once, repeat_with};

#[allow(unused_imports)]
#[cfg(feature = "dbg")]
use dbg::lg;

use crate::fp::fact_build;

fn main() {
    let mut buf = ngtio::with_stdin();
    let n = buf.usize();
    let c = buf.usize();
    let edges = repeat_with(|| [buf.usize() - 1, buf.usize() - 1])
        .take(n - 1)
        .collect::<Vec<_>>();
    let mut degrees = vec![0; n];
    for &[u, v] in &edges {
        degrees[u] += 1;
        degrees[v] += 1;
    }
    let [fact, finv] = fact_build(n + c + 1);
    let cum = once(Fp::new(0))
        .chain((1..=c).map(|c| {
            degrees
                .iter()
                .map(|&d| cal(d, c, &fact, &finv))
                .product::<Fp>()
                * Fp::from(c)
        }))
        .collect::<Vec<_>>();
    let ans = cum
        .iter()
        .enumerate()
        .map(|(i, &x)| {
            x * fact[c]
                * finv[i]
                * finv[c - i]
                * match (c - i) % 2 {
                    0 => Fp::new(1),
                    1 => -Fp::new(1),
                    _ => unreachable!(),
                }
        })
        .sum::<Fp>();
    println!("{}", ans);
}

fn cal(d: usize, c: usize, fact: &[Fp], finv: &[Fp]) -> Fp {
    Fp::from(c).pow((d - 1) as u32)
        - (d / 2 + 1..=d)
            .map(|i| Fp::from(c - 1).pow((d - i) as u32) * fact[d] * finv[i] * finv[d - i])
            .sum::<Fp>()
        + Fp::from(1)
}

// fp {{{
#[allow(dead_code)]
mod fp {
    use std::{
        cmp::PartialEq,
        fmt,
        hash::{Hash, Hasher},
        iter::{successors, Product, Sum},
        marker::PhantomData,
        mem::swap,
        ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
    };
    pub trait Mod: Clone + Copy + Hash {
        const P: u32;
        const K: u32;
        const R2: u32 = ((1_u128 << 64) % Self::P as u128) as _; // 2 ^ 64 mod P
    }
    fn reduce<M: Mod>(x: u64) -> u32 {
        ((x + u64::from(M::K.wrapping_mul(x as u32)) * u64::from(M::P)) >> 32) as u32
    }
    pub fn fact_iter<M: Mod>() -> impl Iterator<Item = Fp<M>> {
        (1..).scan(Fp::new(1), |state, x| {
            let ans = *state;
            *state *= x;
            Some(ans)
        })
    }
    #[allow(clippy::missing_panics_doc)]
    pub fn fact_build<M: Mod>(n: usize) -> [Vec<Fp<M>>; 2] {
        if n == 0 {
            [Vec::new(), Vec::new()]
        } else {
            let fact = fact_iter::<M>().take(n).collect::<Vec<_>>();
            let mut fact_inv = vec![fact.last().unwrap().recip(); n];
            (1..n).rev().for_each(|i| fact_inv[i - 1] = fact_inv[i] * i);
            [fact, fact_inv]
        }
    }
    pub fn binom_iter<M: Mod>() -> impl Iterator<Item = Vec<Fp<M>>> {
        successors(Some(vec![Fp::new(1)]), |last| {
            let mut crr = last.clone();
            crr.push(Fp::new(0));
            crr[1..].iter_mut().zip(last).for_each(|(x, &y)| *x += y);
            Some(crr)
        })
    }
    #[macro_export]
    macro_rules! define_mod {
        ($(($Fp: ident, $Mod: ident, $mod: expr, $k: expr),)*) => {$(
            #[derive(Clone, Debug, Default, Hash, Copy)]
            pub struct $Mod {}
            impl Mod for $Mod {
                const P: u32 = $mod;
                const K: u32 = $k;
            }
            pub type $Fp = Fp<$Mod>;
        )*}
    }
    define_mod! {
        (F998244353, Mod998244353, 998_244_353, 998_244_351),
        (F1000000007, Mod1000000007, 1_000_000_007, 2_226_617_417),
        (F1012924417, Mod1012924417, 1_012_924_417, 1_012_924_415),
        (F924844033, Mod924844033, 924_844_033, 924_844_031),
    }
    #[derive(Clone, Default, Copy)]
    pub struct Fp<M> {
        value: u32,
        __marker: PhantomData<M>,
    }
    impl<M: Mod> Fp<M> {
        pub const P: u32 = M::P;
        pub fn new(value: u32) -> Self {
            Self::from_raw(reduce::<M>(u64::from(value) * u64::from(M::R2)))
        }
        pub fn value(self) -> u32 {
            let x = reduce::<M>(u64::from(self.value));
            if M::P <= x {
                x - M::P
            } else {
                x
            }
        }
        #[allow(clippy::many_single_char_names)]
        pub fn recip(self) -> Self {
            assert_ne!(self, Self::new(0), "0 はだめ。");
            let mut x = M::P as i32;
            let mut y = self.value() as i32;
            let mut u = 0;
            let mut v = 1;
            while y != 0 {
                let q = x / y;
                x -= q * y;
                u -= q * v;
                swap(&mut x, &mut y);
                swap(&mut u, &mut v);
            }
            debug_assert_eq!(x, 1);
            if u < 0 {
                debug_assert_eq!(v, M::P as i32);
                u += v;
            }
            Self::new(u as u32)
        }
        pub fn pow<T: Into<u64>>(self, exp: T) -> Self {
            let mut exp = exp.into();
            if exp == 0 {
                return Self::new(1);
            }
            let mut base = self;
            let mut acc = Self::new(1);
            while 1 < exp {
                if exp & 1 == 1 {
                    acc *= base;
                }
                exp /= 2;
                base *= base;
            }
            acc * base
        }
        fn from_raw(value: u32) -> Self {
            Self {
                value,
                __marker: PhantomData,
            }
        }
    }
    fn simplify(value: i32, p: i32) -> (i32, i32, i32) {
        if value.abs() < 10_000 {
            (value, 1, 0)
        } else {
            let mut q = p.div_euclid(value);
            let mut r = p.rem_euclid(value);
            if value <= 2 * r {
                q += 1;
                r -= value;
            }
            let (num, pden, ppden) = simplify(r, value);
            let den = ppden - q * pden;
            (num, den, pden)
        }
    }
    macro_rules! impl_from_large_int {
        ($($T: ty), *$(,)?) => {$(
            impl<M: Mod> From<$T> for Fp<M> {
                fn from(x: $T) -> Self {
                    Self::new(x.rem_euclid(M::P as _) as u32)
                }
            }
        )*}
    }
    impl_from_large_int! {
        u32, u64, u128, usize,
        i32, i64, i128, isize,
    }
    macro_rules! impl_from_small_int {
        ($($T: ty), *$(,)?) => {$(
            impl<M: Mod> From<$T> for Fp<M> {
                fn from(x: $T) -> Self {
                    Self::new(x as u32)
                }
            }
        )*}
    }
    impl_from_small_int! {
        u8, u16,
        i8, i16,
    }
    impl<M: Mod> PartialEq for Fp<M> {
        fn eq(&self, other: &Self) -> bool {
            fn value<M: Mod>(fp: Fp<M>) -> u32 {
                if fp.value >= M::P {
                    fp.value - M::P
                } else {
                    fp.value
                }
            }
            value(*self) == value(*other)
        }
    }
    impl<M: Mod> Eq for Fp<M> {}
    impl<M: Mod> Hash for Fp<M> {
        fn hash<H: Hasher>(&self, state: &mut H) {
            self.value().hash(state);
        }
    }
    impl<M: Mod, T: Into<Self>> AddAssign<T> for Fp<M> {
        fn add_assign(&mut self, rhs: T) {
            self.value += rhs.into().value;
            if M::P * 2 <= self.value {
                self.value -= M::P * 2;
            }
        }
    }
    impl<M: Mod, T: Into<Self>> SubAssign<T> for Fp<M> {
        fn sub_assign(&mut self, rhs: T) {
            let rhs = rhs.into();
            if self.value < rhs.value {
                self.value += M::P * 2;
            }
            self.value -= rhs.value;
        }
    }
    impl<M: Mod, T: Into<Self>> MulAssign<T> for Fp<M> {
        fn mul_assign(&mut self, rhs: T) {
            self.value = reduce::<M>(u64::from(self.value) * u64::from(rhs.into().value));
        }
    }
    #[allow(clippy::suspicious_op_assign_impl)]
    impl<M: Mod, T: Into<Self>> DivAssign<T> for Fp<M> {
        fn div_assign(&mut self, rhs: T) {
            *self *= rhs.into().recip();
        }
    }
    impl<'a, M: Mod> From<&'a Self> for Fp<M> {
        fn from(x: &Self) -> Self {
            *x
        }
    }
    macro_rules! forward_ops {
        ($(($trait:ident, $method_assign:ident, $method:ident),)*) => {$(
            impl<M: Mod, T: Into<Fp<M>>> $trait<T> for Fp<M> {
                type Output = Self;
                fn $method(mut self, rhs: T) -> Self {
                    self.$method_assign(rhs);
                    self
                }
            }
            impl<'a, M: Mod, T: Into<Fp<M>>> $trait<T> for &'a Fp<M> {
                type Output = Fp<M>;
                fn $method(self, other: T) -> Self::Output {
                    $trait::$method(*self, other)
                }
            }
        )*};
    }
    forward_ops! {
        (Add, add_assign, add),
        (Sub, sub_assign, sub),
        (Mul, mul_assign, mul),
        (Div, div_assign, div),
    }
    impl<M: Mod> Neg for Fp<M> {
        type Output = Self;
        fn neg(self) -> Self {
            Self::from_raw(M::P * 2 - self.value)
        }
    }
    impl<M: Mod> Sum for Fp<M> {
        fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
            iter.fold(Self::new(0), |b, x| b + x)
        }
    }
    impl<M: Mod> Product for Fp<M> {
        fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
            iter.fold(Self::new(1), |b, x| b * x)
        }
    }
    impl<'a, M: Mod> Sum<&'a Self> for Fp<M> {
        fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
            iter.fold(Self::new(0), |b, x| b + x)
        }
    }
    impl<'a, M: Mod> Product<&'a Self> for Fp<M> {
        fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
            iter.fold(Self::new(1), |b, x| b * x)
        }
    }
    impl<M: Mod> fmt::Debug for Fp<M> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
            let (num, den, _) = simplify(self.value() as i32, M::P as i32);
            let (num, den) = match den.signum() {
                1 => (num, den),
                -1 => (-num, -den),
                _ => unreachable!(),
            };
            if den == 1 {
                write!(f, "{}", num)
            } else {
                write!(f, "{}/{}", num, den)
            }
        }
    }
    impl<M: Mod> fmt::Display for Fp<M> {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            self.value().fmt(f)
        }
    }
}
// }}}
// template {{{
#[cfg(not(feature = "dbg"))]
#[allow(unused_macros)]
#[macro_export]
macro_rules! lg {
    ($($expr:expr),*) => {};
}

#[allow(dead_code)]
mod ngtio {

    mod i {
        use std::{
            io::{self, BufRead},
            iter,
        };

        pub use self::{
            multi_token::{Leaf, Parser, ParserTuple, RawTuple, Tuple, VecLen},
            token::{Token, Usize1},
        };

        pub fn with_stdin() -> Tokenizer<io::BufReader<io::Stdin>> {
            io::BufReader::new(io::stdin()).tokenizer()
        }

        pub fn with_str(src: &str) -> Tokenizer<&[u8]> {
            src.as_bytes().tokenizer()
        }

        pub struct Tokenizer<S: BufRead> {
            queue: Vec<String>, // FIXME: String のみにすると速そうです。
            scanner: S,
        }
        macro_rules! prim_method {
            ($name:ident: $T:ty) => {
                pub fn $name(&mut self) -> $T {
                    <$T>::leaf().parse(self)
                }
            };
            ($name:ident) => {
                prim_method!($name: $name);
            };
        }
        macro_rules! prim_methods {
            ($name:ident: $T:ty; $($rest:tt)*) => {
                prim_method!($name:$T);
                prim_methods!($($rest)*);
            };
            ($name:ident; $($rest:tt)*) => {
                prim_method!($name);
                prim_methods!($($rest)*);
            };
            () => ()
        }
        impl<S: BufRead> Tokenizer<S> {
            pub fn token(&mut self) -> String {
                self.load();
                self.queue.pop().expect("入力が終了したのですが。")
            }
            pub fn new(scanner: S) -> Self {
                Self {
                    queue: Vec::new(),
                    scanner,
                }
            }
            fn load(&mut self) {
                while self.queue.is_empty() {
                    let mut s = String::new();
                    let length = self.scanner.read_line(&mut s).unwrap(); // 入力が UTF-8 でないときにエラーだそうです。
                    if length == 0 {
                        break;
                    }
                    self.queue = s.split_whitespace().rev().map(str::to_owned).collect();
                }
            }

            pub fn skip_line(&mut self) {
                assert!(
                    self.queue.is_empty(),
                    "行の途中で呼ばないでいただきたいです。現在のトークンキュー: {:?}",
                    &self.queue
                );
                self.load();
            }

            pub fn end(&mut self) {
                self.load();
                assert!(self.queue.is_empty(), "入力はまだあります！");
            }

            pub fn parse<T: Token>(&mut self) -> T::Output {
                T::parse(&self.token())
            }

            pub fn parse_collect<T: Token, B>(&mut self, n: usize) -> B
            where
                B: iter::FromIterator<T::Output>,
            {
                iter::repeat_with(|| self.parse::<T>()).take(n).collect()
            }

            pub fn tuple<T: RawTuple>(&mut self) -> <T::LeafTuple as Parser>::Output {
                T::leaf_tuple().parse(self)
            }

            pub fn vec<T: Token>(&mut self, len: usize) -> Vec<T::Output> {
                T::leaf().vec(len).parse(self)
            }

            pub fn vec_tuple<T: RawTuple>(
                &mut self,
                len: usize,
            ) -> Vec<<T::LeafTuple as Parser>::Output> {
                T::leaf_tuple().vec(len).parse(self)
            }

            pub fn vec2<T: Token>(&mut self, height: usize, width: usize) -> Vec<Vec<T::Output>> {
                T::leaf().vec(width).vec(height).parse(self)
            }

            pub fn vec2_tuple<T>(
                &mut self,
                height: usize,
                width: usize,
            ) -> Vec<Vec<<T::LeafTuple as Parser>::Output>>
            where
                T: RawTuple,
            {
                T::leaf_tuple().vec(width).vec(height).parse(self)
            }
            prim_methods! {
                u8; u16; u32; u64; u128; usize;
                i8; i16; i32; i64; i128; isize;
                f32; f64;
                char; string: String;
            }
        }

        mod token {
            use super::multi_token::Leaf;
            use std::{any, fmt, marker, str};

            pub trait Token: Sized {
                type Output;
                fn parse(s: &str) -> Self::Output;
                fn leaf() -> Leaf<Self> {
                    Leaf(marker::PhantomData)
                }
            }

            impl<T> Token for T
            where
                T: str::FromStr,
                <T as str::FromStr>::Err: fmt::Debug,
            {
                type Output = T;
                fn parse(s: &str) -> Self::Output {
                    s.parse().unwrap_or_else(|_| {
                        panic!("Parse error!: ({}: {})", s, any::type_name::<T>(),)
                    })
                }
            }

            pub struct Usize1 {}
            impl Token for Usize1 {
                type Output = usize;
                fn parse(s: &str) -> Self::Output {
                    usize::parse(s)
                        .checked_sub(1)
                        .expect("Parse error! (Zero substruction error of Usize1)")
                }
            }
        }

        mod multi_token {
            use super::{Token, Tokenizer};
            use std::{io::BufRead, iter, marker};

            pub trait Parser: Sized {
                type Output;
                fn parse<S: BufRead>(&self, server: &mut Tokenizer<S>) -> Self::Output;
                fn vec(self, len: usize) -> VecLen<Self> {
                    VecLen { len, elem: self }
                }
            }
            pub struct Leaf<T>(pub(super) marker::PhantomData<T>);
            impl<T: Token> Parser for Leaf<T> {
                type Output = T::Output;
                fn parse<S: BufRead>(&self, server: &mut Tokenizer<S>) -> T::Output {
                    server.parse::<T>()
                }
            }

            pub struct VecLen<T> {
                pub len: usize,
                pub elem: T,
            }
            impl<T: Parser> Parser for VecLen<T> {
                type Output = Vec<T::Output>;
                fn parse<S: BufRead>(&self, server: &mut Tokenizer<S>) -> Self::Output {
                    iter::repeat_with(|| self.elem.parse(server))
                        .take(self.len)
                        .collect()
                }
            }

            pub trait RawTuple {
                type LeafTuple: Parser;
                fn leaf_tuple() -> Self::LeafTuple;
            }
            pub trait ParserTuple {
                type Tuple: Parser;
                fn tuple(self) -> Self::Tuple;
            }
            pub struct Tuple<T>(pub T);
            macro_rules! impl_tuple {
                ($($t:ident: $T:ident),*) => {
                    impl<$($T),*> Parser for Tuple<($($T,)*)>
                        where
                            $($T: Parser,)*
                            {
                                type Output = ($($T::Output,)*);
#[allow(unused_variables)]
                                fn parse<S: BufRead >(&self, server: &mut Tokenizer<S>) -> Self::Output {
                                    match self {
                                        Tuple(($($t,)*)) => {
                                            ($($t.parse(server),)*)
                                        }
                                    }
                                }
                            }
                    impl<$($T: Token),*> RawTuple for ($($T,)*) {
                        type LeafTuple = Tuple<($(Leaf<$T>,)*)>;
                        fn leaf_tuple() -> Self::LeafTuple {
                            Tuple(($($T::leaf(),)*))
                        }
                    }
                    impl<$($T: Parser),*> ParserTuple for ($($T,)*) {
                        type Tuple = Tuple<($($T,)*)>;
                        fn tuple(self) -> Self::Tuple {
                            Tuple(self)
                        }
                    }
                };
            }
            impl_tuple!();
            impl_tuple!(t1: T1);
            impl_tuple!(t1: T1, t2: T2);
            impl_tuple!(t1: T1, t2: T2, t3: T3);
            impl_tuple!(t1: T1, t2: T2, t3: T3, t4: T4);
            impl_tuple!(t1: T1, t2: T2, t3: T3, t4: T4, t5: T5);
            impl_tuple!(t1: T1, t2: T2, t3: T3, t4: T4, t5: T5, t6: T6);
            impl_tuple!(t1: T1, t2: T2, t3: T3, t4: T4, t5: T5, t6: T6, t7: T7);
            impl_tuple!(
                t1: T1,
                t2: T2,
                t3: T3,
                t4: T4,
                t5: T5,
                t6: T6,
                t7: T7,
                t8: T8
            );
        }

        trait Scanner: BufRead + Sized {
            fn tokenizer(self) -> Tokenizer<Self> {
                Tokenizer::new(self)
            }
        }
        impl<R: BufRead> Scanner for R {}
    }

    pub use self::i::{with_stdin, with_str};

    pub mod prelude {
        pub use super::i::{Parser, ParserTuple, RawTuple, Token, Usize1};
    }
}
// }}}

yukicoder

結果

テストケース

コンパイルメッセージ

ソースコード