#1133871 (Rust) No.3328 岩井ツリーグラフ

• 問題: No.3328 岩井ツリーグラフ
• ユーザー: elphe
• 提出日時: 2025-11-15 14:04:34
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 30 ms / 2,000 ms
• コード長: 7,480 bytes
• コンパイル時間: 11,058 ms
• コンパイル使用メモリ: 398,380 KB
• 実行使用メモリ: 7,720 KB
• 最終ジャッジ日時: 2025-11-15 14:04:50
• 合計ジャッジ時間: 13,666 ms
• ジャッジサーバーID （参考情報）: judge2 / judge4

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

ソースコード

use proconio::{fastout, input};

mod mylib {
    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
    pub struct ModInt<const MOD: u32> {
        value: u32,
    }

    use std::ops;
    #[allow(dead_code)]
    impl<const MOD: u32> ModInt<MOD> {
        const fn is_mod_prime() -> bool {
            let mut i = 2u32;
            while i * i <= MOD {
                if MOD % i == 0 {
                    return false;
                }
                i += 1;
            }
            true
        }

        pub const fn pow(self, mut a: u32) -> ModInt<MOD> {
            let mut ans: u32 = 1;
            let mut mul: u32 = self.value;
            while a > 0 {
                if (a & 1) == 1 {
                    ans = (ans as u64 * mul as u64 % MOD as u64) as u32;
                }
                mul = (mul as u64 * mul as u64 % MOD as u64) as u32;
                a /= 2;
            }
            Self { value: ans }
        }

        pub const fn add(self, other: Self) -> Self {
            if other.value >= MOD - self.value {
                Self {
                    value: other.value - (MOD - self.value),
                }
            } else {
                Self {
                    value: self.value + other.value,
                }
            }
        }

        pub const fn sub(self, other: Self) -> Self {
            if self.value >= other.value {
                Self {
                    value: self.value - other.value,
                }
            } else {
                Self {
                    value: self.value + (MOD - other.value),
                }
            }
        }

        pub const fn mul(self, other: Self) -> Self {
            if const { MOD > u16::MAX as u32 } {
                Self {
                    value: (self.value as u64 * other.value as u64 % MOD as u64) as u32,
                }
            } else {
                Self {
                    value: self.value * other.value % MOD,
                }
            }
        }

        pub const fn div(self, other: Self) -> Self {
            const {
                assert!(Self::is_mod_prime());
            }
            Self {
                value: (self.value as u64 * other.pow(MOD - 2).value as u64 % MOD as u64) as u32,
            }
        }

        pub const fn permutations_mod(n: u32, r: u32) -> Self {
            let mut ans: u32 = 1;
            let mut i: u32 = 0;
            while i < r {
                ans = (ans as u64 * (n - i) as u64 % MOD as u64) as u32;
                i += 1;
            }
            Self { value: ans }
        }

        pub const fn combinations_mod(n: u32, r: u32) -> Self {
            if r * 2 > n {
                Self::combinations_mod(n, n - r)
            } else {
                Self::permutations_mod(n, r).mul(Self::pow(Self::permutations_mod(r, r), MOD - 2))
            }
        }

        pub const fn from_direct(n: u32) -> Self {
            Self { value: n }
        }
    }

    impl<const MOD: u32> From<u32> for ModInt<MOD> {
        fn from(n: u32) -> Self {
            Self { value: n % MOD }
        }
    }

    impl<const MOD: u32> From<u64> for ModInt<MOD> {
        fn from(n: u64) -> Self {
            Self {
                value: (n % MOD as u64) as u32,
            }
        }
    }

    impl<const MOD: u32> Into<u32> for ModInt<MOD> {
        fn into(self) -> u32 {
            self.value
        }
    }

    impl<const MOD: u32> Into<u64> for ModInt<MOD> {
        fn into(self) -> u64 {
            self.value as u64
        }
    }

    use std::fmt;
    impl<const MOD: u32> fmt::Display for ModInt<MOD> {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            write!(f, "{}", self.value)
        }
    }

    impl<const MOD: u32> ops::Add for ModInt<MOD> {
        type Output = Self;
        fn add(self, other: ModInt<MOD>) -> <Self as ops::Add<ModInt<MOD>>>::Output {
            if other.value >= MOD - self.value {
                Self {
                    value: other.value - (MOD - self.value),
                }
            } else {
                Self {
                    value: self.value + other.value,
                }
            }
        }
    }

    impl<const MOD: u32> ops::Sub for ModInt<MOD> {
        type Output = Self;
        fn sub(self, other: ModInt<MOD>) -> <Self as ops::Sub<ModInt<MOD>>>::Output {
            if self.value >= other.value {
                Self {
                    value: self.value - other.value,
                }
            } else {
                Self {
                    value: self.value + (MOD - other.value),
                }
            }
        }
    }

    impl<const MOD: u32> ops::Mul for ModInt<MOD> {
        type Output = Self;
        fn mul(self, other: ModInt<MOD>) -> <Self as ops::Mul<ModInt<MOD>>>::Output {
            if const { MOD > u16::MAX as u32 } {
                Self {
                    value: (self.value as u64 * other.value as u64 % MOD as u64) as u32,
                }
            } else {
                Self {
                    value: self.value * other.value % MOD,
                }
            }
        }
    }

    impl<const MOD: u32> ops::Div for ModInt<MOD> {
        type Output = Self;
        fn div(self, other: ModInt<MOD>) -> <Self as ops::Div<ModInt<MOD>>>::Output {
            const {
                assert!(Self::is_mod_prime());
            }
            Self {
                value: (self.value as u64 * other.pow(MOD - 2).value as u64 % MOD as u64) as u32,
            }
        }
    }

    impl<const MOD: u32> ops::AddAssign for ModInt<MOD> {
        fn add_assign(&mut self, other: Self) {
            if self.value >= MOD - other.value {
                self.value -= MOD - other.value;
            } else {
                self.value += other.value;
            }
        }
    }

    impl<const MOD: u32> ops::SubAssign for ModInt<MOD> {
        fn sub_assign(&mut self, other: Self) {
            if self.value >= other.value {
                self.value -= other.value;
            } else {
                self.value += MOD - other.value;
            }
        }
    }

    impl<const MOD: u32> ops::MulAssign for ModInt<MOD> {
        fn mul_assign(&mut self, other: Self) {
            if const { MOD > u16::MAX as u32 } {
                self.value = (self.value as u64 * other.value as u64 % MOD as u64) as u32;
            } else {
                self.value = self.value * other.value % MOD;
            }
        }
    }

    impl<const MOD: u32> ops::DivAssign for ModInt<MOD> {
        fn div_assign(&mut self, other: Self) {
            const {
                assert!(Self::is_mod_prime());
            }
            self.value = (self.value as u64 * other.pow(MOD - 2).value as u64 % MOD as u64) as u32;
        }
    }
}

#[fastout]
fn main() {
    input! {
        y: [u32],
    }

    println!("{}", output(solve(y)));
}

fn solve(y: Vec<u32>) -> u32 {
    const MOD: u32 = 998244353;
	const ONE_OF_THREE: mylib::ModInt<MOD> = mylib::ModInt::<MOD>::from_direct(1).div(mylib::ModInt::<MOD>::from_direct(3));
    let sum_y = mylib::ModInt::<MOD>::from(y.iter().fold(0, |x, y| x + *y as u64));
    let mut ans = mylib::ModInt::<MOD>::from_direct(0u32);

    for y in y {
        ans += mylib::ModInt::<MOD>::from(y as u64 * (y + 1) as u64 / 2) * (sum_y - mylib::ModInt::<MOD>::from(2 * y - 2) * ONE_OF_THREE);
    }
    ans.into()
}

fn output(ans: u32) -> u32 {
    ans
}