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

• 問題: No.3328 岩井ツリーグラフ
• ユーザー: yiwiy9
• 提出日時: 2026-09-05 02:16:02
• 言語: Rust (1.97.1 + proconio + num + itertools)
• 結果: AC
• 実行時間: 142 ms / 2,000 ms
• コード長: 7,172 bytes
• コンパイル時間: 671 ms
• コンパイル使用メモリ: 209,832 KB
• 実行使用メモリ: 9,788 KB
• 最終ジャッジ日時: 2026-09-05 02:16:08
• 合計ジャッジ時間: 5,240 ms
• ジャッジサーバーID （参考情報）: judge4_0 / judge2_0

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

コンパイルメッセージ

warning: associated function `modulus` is never used
  --> src/main.rs:68:8
   |
67 | impl Mint {
   | --------- associated function in this implementation
68 |     fn modulus() -> u32 {
   |        ^^^^^^^
   |
   = note: `#[warn(dead_code)]` (part of `#[warn(unused)]`) on by default

ソースコード

use proconio::input;

fn main() {
    input! {
        x: usize,
        y: [Mint; x],
    }

    let mut ans = Mint::new(0);
    let mut sum_y = Mint::new(0);
    let mut sum_depth = Mint::new(0);
    for &y_i in &y {
        ans += y_i * (y_i + 1) * (Mint::new(2) * y_i + 1) / Mint::new(2) / Mint::new(6)
            + y_i * (y_i + 1) / Mint::new(2) / Mint::new(2);
        ans += sum_depth * y_i + sum_y * (y_i * (y_i + 1) / Mint::new(2));
        sum_y += y_i;
        sum_depth += y_i * (y_i + 1) / Mint::new(2);
    }

    println!("{}", ans.val());
}

// From here: a compact subset of ac-library-rs's static modint API.
use std::{
    fmt::{self, Debug, Display, Formatter},
    iter::{Product, Sum},
    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
    str::FromStr,
};

const MOD: u32 = 998_244_353;

trait RemEuclidU32 {
    fn rem_euclid_u32(self, modulus: u32) -> u32;
}

macro_rules! impl_rem_euclid_u32_for_signed {
    ($($ty:ty),*) => {
        $(
            impl RemEuclidU32 for $ty {
                fn rem_euclid_u32(self, modulus: u32) -> u32 {
                    (self as i128).rem_euclid(modulus as i128) as u32
                }
            }
        )*
    };
}

macro_rules! impl_rem_euclid_u32_for_unsigned {
    ($($ty:ty),*) => {
        $(
            impl RemEuclidU32 for $ty {
                fn rem_euclid_u32(self, modulus: u32) -> u32 {
                    (self as u128 % modulus as u128) as u32
                }
            }
        )*
    };
}

impl_rem_euclid_u32_for_signed!(i8, i16, i32, i64, i128, isize);
impl_rem_euclid_u32_for_unsigned!(u8, u16, u32, u64, u128, usize);

#[derive(Clone, Copy, Default, Eq, PartialEq, Hash)]
struct Mint(u32);

impl Mint {
    fn modulus() -> u32 {
        MOD
    }

    fn new<T: RemEuclidU32>(value: T) -> Self {
        Self(value.rem_euclid_u32(MOD))
    }

    fn raw(value: u32) -> Self {
        debug_assert!(value < MOD);
        Self(value)
    }

    fn val(self) -> u32 {
        self.0
    }

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

    fn inv(self) -> Self {
        assert!(self.0 != 0, "attempt to divide by zero");
        self.pow((MOD - 2) as u64)
    }
}

impl FromStr for Mint {
    type Err = <i128 as FromStr>::Err;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        s.parse::<i128>().map(Self::new)
    }
}

impl Display for Mint {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        Display::fmt(&self.0, f)
    }
}

impl Debug for Mint {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        Display::fmt(self, f)
    }
}

impl Add for Mint {
    type Output = Self;

    fn add(self, rhs: Self) -> Self {
        let value = self.0 + rhs.0;
        Self::raw(if value >= MOD { value - MOD } else { value })
    }
}

impl Sub for Mint {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self {
        Self::raw(if self.0 >= rhs.0 {
            self.0 - rhs.0
        } else {
            self.0 + MOD - rhs.0
        })
    }
}

impl Mul for Mint {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self {
        Self::raw((self.0 as u64 * rhs.0 as u64 % MOD as u64) as u32)
    }
}

impl Div for Mint {
    type Output = Self;

    fn div(self, rhs: Self) -> Self {
        self * rhs.inv()
    }
}

impl Neg for Mint {
    type Output = Self;

    fn neg(self) -> Self {
        if self.0 == 0 {
            self
        } else {
            Self::raw(MOD - self.0)
        }
    }
}

impl AddAssign for Mint {
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl SubAssign for Mint {
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl MulAssign for Mint {
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs;
    }
}

impl DivAssign for Mint {
    fn div_assign(&mut self, rhs: Self) {
        *self = *self / rhs;
    }
}

macro_rules! impl_ops_with_integer {
    ($($ty:ty),*) => {
        $(
            impl Add<$ty> for Mint {
                type Output = Self;

                fn add(self, rhs: $ty) -> Self {
                    self + Self::new(rhs)
                }
            }

            impl Sub<$ty> for Mint {
                type Output = Self;

                fn sub(self, rhs: $ty) -> Self {
                    self - Self::new(rhs)
                }
            }

            impl Mul<$ty> for Mint {
                type Output = Self;

                fn mul(self, rhs: $ty) -> Self {
                    self * Self::new(rhs)
                }
            }

            impl Div<$ty> for Mint {
                type Output = Self;

                fn div(self, rhs: $ty) -> Self {
                    self / Self::new(rhs)
                }
            }

            impl AddAssign<$ty> for Mint {
                fn add_assign(&mut self, rhs: $ty) {
                    *self += Self::new(rhs);
                }
            }

            impl SubAssign<$ty> for Mint {
                fn sub_assign(&mut self, rhs: $ty) {
                    *self -= Self::new(rhs);
                }
            }

            impl MulAssign<$ty> for Mint {
                fn mul_assign(&mut self, rhs: $ty) {
                    *self *= Self::new(rhs);
                }
            }

            impl DivAssign<$ty> for Mint {
                fn div_assign(&mut self, rhs: $ty) {
                    *self /= Self::new(rhs);
                }
            }

            impl Add<Mint> for $ty {
                type Output = Mint;

                fn add(self, rhs: Mint) -> Mint {
                    Mint::new(self) + rhs
                }
            }

            impl Sub<Mint> for $ty {
                type Output = Mint;

                fn sub(self, rhs: Mint) -> Mint {
                    Mint::new(self) - rhs
                }
            }

            impl Mul<Mint> for $ty {
                type Output = Mint;

                fn mul(self, rhs: Mint) -> Mint {
                    Mint::new(self) * rhs
                }
            }

            impl Div<Mint> for $ty {
                type Output = Mint;

                fn div(self, rhs: Mint) -> Mint {
                    Mint::new(self) / rhs
                }
            }
        )*
    };
}

impl_ops_with_integer!(i8, i16, i32, i64, i128, isize);
impl_ops_with_integer!(u8, u16, u32, u64, u128, usize);

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

impl<'a> Sum<&'a Self> for Mint {
    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
        iter.copied().sum()
    }
}

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

impl<'a> Product<&'a Self> for Mint {
    fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
        iter.copied().product()
    }
}