結果

問題 No.235 めぐるはめぐる (5)
ユーザー HaarHaar
提出日時 2021-12-04 09:11:44
言語 Rust
(1.77.0)
結果
AC  
実行時間 1,426 ms / 10,000 ms
コード長 25,080 bytes
コンパイル時間 5,231 ms
コンパイル使用メモリ 179,128 KB
実行使用メモリ 68,064 KB
最終ジャッジ日時 2023-09-20 19:06:24
合計ジャッジ時間 11,173 ms
ジャッジサーバーID
(参考情報)
judge14 / judge11
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1,426 ms
65,388 KB
testcase_01 AC 887 ms
68,064 KB
testcase_02 AC 1,313 ms
65,972 KB
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: field `size` is never read
   --> Main.rs:633:13
    |
632 |         pub struct HLD {
    |                    --- field in this struct
633 |             size: usize,
    |             ^^^^
    |
    = note: `HLD` has derived impls for the traits `Debug` and `Clone`, but these are intentionally ignored during dead code analysis
    = note: `#[warn(dead_code)]` on by default

warning: 1 warning emitted

ソースコード

diff #

#![allow(unused_imports)]
use std::io::{Read, Write};

#[allow(unused_macros)]
macro_rules! get {
    ( $in:ident, [$a:tt; $num:expr] ) => {
        {
            let n = $num;
            (0 .. n).map(|_| get!($in, $a)).collect::<Vec<_>>()
        }
    };

    ( $in:ident, ($($type:ty),*) ) => {
        ($(get!($in, $type)),*)
    };

    ( $in:ident, $type:ty ) => {
        {
            let token = $in.next().unwrap();

            token.parse::<$type>().expect(
                format!("cannot convert \"{}\" into {}", token, stringify!($type)).as_str()
            )
        }
    };
}

macro_rules! input {
    ( @inner $in:ident, mut $name:ident : $type:tt ) => {
        let mut $name = get!($in, $type);
    };

    ( @inner $in:ident, $name:ident : $type:tt ) => {
        let $name = get!($in, $type);
    };

    ( $in:ident, $($($names:ident)* : $type:tt),* ) => {
        $(
            input!(@inner $in, $($names)* : $type);
        )*
    }
}

#[allow(unused_macros)]
macro_rules! io {
    ( $in:ident, $out:ident ) => {
        let mut s = String::new();
        std::io::stdin().read_to_string(&mut s).unwrap();
        let mut $in = s.split_ascii_whitespace();

        let $out = std::io::stdout();
        let mut $out = std::io::BufWriter::new($out.lock());
    };
}

pub mod main {
    use super::*;
    use haar_lib::{
        ds::lazy_segtree_coeff::*,
        math::ff::modint::*,
        modulo,
        sort_with,
        //chmin, chmax,
        //mul_vec,
        tree::{hld::*, *},
    };

    #[derive(Clone, Default)]
    pub struct Problem {/* write variables here */}

    modulo!(M, 1000000007);
    type Mint = ModInt<M>;

    impl Problem {
        pub fn main(&mut self) -> Result<(), Box<dyn std::error::Error>> {
            io!(cin, cout);

            input!(cin, n: usize, mut s: [Mint; n], mut c: [Mint; n]);

            let mut tree = Tree::new(n);
            for _ in 0..n - 1 {
                input!(cin, u: usize, v: usize);
                tree.add_undirected(Some((u - 1, v - 1, ())));
            }

            let hld = HLD::new(tree, 0);

            sort_with!(|&i, &j| hld.get_id(i).cmp(&hld.get_id(j)), s, c);

            let mut seg = LazySegmentTreeCoeff::new(n, c);
            seg.init_with_vec(s);

            input!(cin, q: usize);

            for _ in 0..q {
                input!(cin, ty: usize);

                match ty {
                    0 => {
                        input!(cin, x: usize, y: usize, z: Mint);
                        for (l, r) in hld.path_query_vertex(x - 1, y - 1) {
                            seg.update(l..r, z);
                        }
                    }
                    1 => {
                        input!(cin, x: usize, y: usize);
                        let mut ans = Mint::from(0);
                        for (l, r) in hld.path_query_vertex(x - 1, y - 1) {
                            ans += seg.fold(l..r);
                        }
                        writeln!(cout, "{}", ans)?;
                    }
                    _ => unreachable!(),
                }
            }

            Ok(())
        }
        /* write functions here */
    }
}

fn main() {
    main::Problem::default().main().unwrap();
}
use crate as haar_lib;
pub mod algebra {
    pub mod one_zero {
        pub trait Zero {
            type Output;
            fn zero() -> Self::Output;
        }

        pub trait One {
            type Output;
            fn one() -> Self::Output;
        }

        macro_rules! impl_one_zero
{
    ($($t:ty),*) => {
        $(
            impl Zero for $t {
                type Output = $t;
                fn zero() -> Self::Output { 0 as $t }
            }
            impl One for $t {
                type Output = $t;
                fn one() -> Self::Output { 1 as $t }
            }
        )*
    }
}

        impl_one_zero!(u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize, f32, f64);
    }
}
pub mod ds {
    pub mod traits {
        pub trait Foldable<Idx> {
            type Output;
            fn fold(&self, range: Idx) -> Self::Output;
        }

        pub trait FoldableMut<Idx> {
            type Output;
            fn fold(&mut self, range: Idx) -> Self::Output;
        }

        pub trait Assignable<Idx> {
            type Value;
            fn assign(&mut self, i: Idx, value: Self::Value);
        }

        pub trait Updatable<Idx> {
            type Value;
            fn update(&mut self, i: Idx, value: Self::Value);
        }

        pub trait IndexableMut<Idx> {
            type Output;
            fn get(&mut self, i: Idx) -> Self::Output;
        }
    }

    pub mod lazy_segtree_coeff {
        pub use crate::ds::traits::{FoldableMut, Updatable};

        use std::ops::{Add, Mul, Range};

        pub struct LazySegmentTreeCoeff<T> {
            size: usize,
            data: Vec<T>,
            lazy: Vec<T>,
            coeff: Vec<T>,
        }

        impl<T> LazySegmentTreeCoeff<T>
        where
            T: Copy + Default + Add<Output = T> + Mul<Output = T> + PartialEq,
        {
            pub fn new(n: usize, coefficients: Vec<T>) -> Self {
                let size = n.next_power_of_two() * 2;
                let mut coeff = vec![T::default(); size];
                for i in 0..coefficients.len() {
                    coeff[i + size / 2] = coefficients[i];
                }
                for i in (1..size / 2).rev() {
                    coeff[i] = coeff[i << 1] + coeff[i << 1 | 1];
                }
                Self {
                    size,
                    data: vec![T::default(); size],
                    lazy: vec![T::default(); size],
                    coeff,
                }
            }

            pub fn init_with_vec(&mut self, value: Vec<T>) {
                self.data = vec![T::default(); self.size];
                self.lazy = vec![T::default(); self.size];
                for (i, x) in value.into_iter().enumerate() {
                    self.data[self.size / 2 + i] = x;
                }
                for i in (1..self.size / 2).rev() {
                    self.data[i] = self.data[i << 1] + self.data[i << 1 | 1];
                }
            }

            fn propagate(&mut self, i: usize) {
                if self.lazy[i] != T::default() {
                    if i < self.size / 2 {
                        self.lazy[i << 1] = self.lazy[i] + self.lazy[i << 1];
                        self.lazy[i << 1 | 1] = self.lazy[i] + self.lazy[i << 1 | 1];
                    }
                    self.data[i] = self.data[i] + self.lazy[i] * self.coeff[i];
                    self.lazy[i] = T::default();
                }
            }

            fn update_internal(
                &mut self,
                i: usize,
                l: usize,
                r: usize,
                s: usize,
                t: usize,
                value: T,
            ) -> T {
                self.propagate(i);
                if r <= s || t <= l {
                    return self.data[i];
                }
                if s <= l && r <= t {
                    self.lazy[i] = self.lazy[i] + value;
                    self.propagate(i);
                    return self.data[i];
                }
                self.data[i] = self.update_internal(i << 1, l, (l + r) / 2, s, t, value)
                    + self.update_internal(i << 1 | 1, (l + r) / 2, r, s, t, value);
                self.data[i]
            }

            fn get_internal(&mut self, i: usize, l: usize, r: usize, x: usize, y: usize) -> T {
                self.propagate(i);
                if r <= x || y <= l {
                    return T::default();
                }
                if x <= l && r <= y {
                    return self.data[i];
                }
                self.get_internal(i << 1, l, (l + r) / 2, x, y)
                    + self.get_internal(i << 1 | 1, (l + r) / 2, r, x, y)
            }
        }

        impl<T> Updatable<Range<usize>> for LazySegmentTreeCoeff<T>
        where
            T: Copy + Default + Add<Output = T> + Mul<Output = T> + PartialEq,
        {
            type Value = T;

            fn update(&mut self, Range { start, end }: Range<usize>, value: T) {
                self.update_internal(1, 0, self.size / 2, start, end, value);
            }
        }

        impl<T> FoldableMut<Range<usize>> for LazySegmentTreeCoeff<T>
        where
            T: Copy + Default + Add<Output = T> + Mul<Output = T> + PartialEq,
        {
            type Output = T;

            fn fold(&mut self, Range { start, end }: Range<usize>) -> T {
                self.get_internal(1, 0, self.size / 2, start, end)
            }
        }
    }
}
pub mod macros {
    pub mod modulo {
        #[macro_export]
        macro_rules! modulo {
            ($name:ident, $m:expr) => {
                #[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
                struct $name {}
                impl Modulo for $name {
                    #[inline]
                    fn value() -> u32 {
                        $m
                    }
                }
            };
        }
    }

    pub mod sort_with {
        #[macro_export]
        macro_rules! sort_with
{
    ($cmp:expr, $($a:expr),+) => {
        {
            let n = vec![$($a.len()),+].into_iter().min().unwrap();
            let mut ord = (0..n).collect::<Vec<usize>>();
            ord.sort_by($cmp);
            let mut check = vec![false; n];
            for i in 0..n {
                if !check[i] {
                    check[i] = true;
                    let mut j = i;
                    while ord[j] != i {
                        {
                            $(
                                $a.swap(j, ord[j]);
                            )+
                        }
                        j = ord[j];
                        check[j] = true;
                    }
                }
            }
        }
    }
}
    }
}
pub mod math {
    pub mod ff {
        pub mod modint {
            pub use crate::{
                algebra::one_zero::{One, Zero},
                math::ff::traits::{Frac, Inv, Pow, FF},
            };

            use std::{
                fmt,
                fmt::{Debug, Display, Formatter},
                iter::Sum,
                marker::PhantomData,
                num::ParseIntError,
                ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
                str::FromStr,
            };

            pub trait Modulo {
                fn value() -> u32;
            }

            #[derive(Copy, Clone, PartialEq, Default)]
            pub struct ModInt<M> {
                value: u32,
                phantom: PhantomData<M>,
            }

            impl<M: Modulo + Copy + PartialEq + Default> FF for ModInt<M> {}

            impl<M: Modulo> ModInt<M> {
                pub fn new(n: u32) -> Self {
                    Self {
                        value: if n < M::value() { n } else { n % M::value() },
                        phantom: PhantomData,
                    }
                }

                #[inline]
                fn new_unchecked(value: u32) -> Self {
                    Self {
                        value,
                        phantom: PhantomData,
                    }
                }

                #[inline]
                fn add_internal(self, other: Self) -> Self {
                    let a = self.value + other.value;
                    Self::new_unchecked(if a < M::value() { a } else { a - M::value() })
                }

                #[inline]
                fn sub_internal(self, other: Self) -> Self {
                    let a = if self.value < other.value {
                        self.value + M::value() - other.value
                    } else {
                        self.value - other.value
                    };
                    Self::new_unchecked(a)
                }

                #[inline]
                fn mul_internal(self, other: Self) -> Self {
                    let a = self.value as u64 * other.value as u64;
                    Self::new_unchecked(if a < M::value() as u64 {
                        a as u32
                    } else {
                        (a % M::value() as u64) as u32
                    })
                }

                #[inline]
                fn div_internal(self, other: Self) -> Self {
                    self * other.inv_internal()
                }

                #[inline]
                fn inv_internal(self) -> Self {
                    self.pow_internal(M::value() as u64 - 2)
                }

                #[inline]
                fn pow_internal(self, mut p: u64) -> Self {
                    let mut ret: u64 = 1;
                    let mut a = self.value as u64;
                    while p > 0 {
                        if (p & 1) != 0 {
                            ret *= a;
                            ret %= M::value() as u64;
                        }
                        a *= a;
                        a %= M::value() as u64;
                        p >>= 1;
                    }
                    Self::new_unchecked(ret as u32)
                }
            }

            impl<M: Modulo> Pow for ModInt<M> {
                type Output = Self;

                fn pow(self, p: u64) -> Self {
                    self.pow_internal(p)
                }
            }

            impl<M: Modulo> Inv for ModInt<M> {
                type Output = Self;

                fn inv(self) -> Self {
                    self.inv_internal()
                }
            }

            impl<M: Modulo> Frac for ModInt<M> {
                type Output = Self;

                fn frac(numerator: i64, denominator: i64) -> Self {
                    Self::from(numerator) * Self::from(denominator).inv()
                }
            }

            impl<M: Modulo> Display for ModInt<M> {
                fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
                    write!(f, "{}", self.value)
                }
            }

            impl<M: Modulo> Debug for ModInt<M> {
                fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
                    write!(f, "{} (mod {})", self.value, M::value())
                }
            }

            macro_rules! modint_from_int
{
    ( $($t:ty),* ) => {
        $(
            impl<M: Modulo> From<$t> for ModInt<M> {
                fn from(from: $t) -> Self {
                    let mut value = ((from % M::value() as $t) + M::value() as $t) as u32;
                    if value >= M::value() {
                        value -= M::value();
                    }
                    Self::new_unchecked(value)
                }
            }
        )*
    }
}

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

            impl<M> From<ModInt<M>> for u32 {
                fn from(from: ModInt<M>) -> Self {
                    from.value
                }
            }

            macro_rules! impl_modint_arith
{
    ($tr:ident, $f:ident, $fi:ident, $tr_a:ident, $f_a:ident, $op:tt) => {
        impl<M: Modulo> $tr for ModInt<M> {
            type Output = Self;
            #[inline]
            fn $f(self, other: Self) -> Self {
                self.$fi(other)
            }
        }
        impl<M: Modulo + Copy> $tr_a for ModInt<M> {
            #[inline]
            fn $f_a(&mut self, other: Self) {
                *self = *self $op other;
            }
        }
    }
}

            impl_modint_arith ! (Add , add , add_internal , AddAssign , add_assign , +);

            impl_modint_arith ! (Sub , sub , sub_internal , SubAssign , sub_assign , -);

            impl_modint_arith ! (Mul , mul , mul_internal , MulAssign , mul_assign , *);

            impl_modint_arith ! (Div , div , div_internal , DivAssign , div_assign , /);

            impl<M: Modulo> Neg for ModInt<M> {
                type Output = Self;

                fn neg(self) -> Self {
                    Self::new_unchecked(if self.value == 0 {
                        0
                    } else {
                        M::value() - self.value
                    })
                }
            }

            impl<M: Modulo> FromStr for ModInt<M> {
                type Err = ParseIntError;

                fn from_str(s: &str) -> Result<Self, Self::Err> {
                    let x = s.parse::<u32>()?;
                    Ok(Self::from(x))
                }
            }

            impl<M: Modulo> Sum for ModInt<M> {
                fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
                    iter.fold(Self::new_unchecked(0), |a, b| a + b)
                }
            }

            impl<M: Modulo> Zero for ModInt<M> {
                type Output = Self;

                #[inline]
                fn zero() -> Self::Output {
                    Self::new_unchecked(0)
                }
            }

            impl<M: Modulo> One for ModInt<M> {
                type Output = Self;

                #[inline]
                fn one() -> Self::Output {
                    Self::new_unchecked(1)
                }
            }
        }

        pub mod traits {
            use std::{
                iter::Sum,
                ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
            };

            pub trait Pow {
                type Output;
                fn pow(self, p: u64) -> Self::Output;
            }

            pub trait Inv {
                type Output;
                fn inv(self) -> Self::Output;
            }

            pub trait Frac {
                type Output;
                fn frac(_: i64, _: i64) -> Self::Output;
            }

            pub trait FF:
                Pow<Output = Self>
                + Inv<Output = Self>
                + Frac<Output = Self>
                + Add<Output = Self>
                + AddAssign
                + Sub<Output = Self>
                + SubAssign
                + Mul<Output = Self>
                + MulAssign
                + Div<Output = Self>
                + DivAssign
                + Neg<Output = Self>
                + Sum
                + Copy
                + Clone
                + PartialEq
                + Default
                + Sized
            {
            }
        }
    }
}
pub mod tree {
    pub mod hld {
        use crate::tree::*;

        use std::cmp::max;

        #[derive(Clone, Debug)]
        pub struct HLD {
            size: usize,
            par: Vec<Option<usize>>,
            head: Vec<usize>,
            id: Vec<usize>,
            rid: Vec<usize>,
            next: Vec<Option<usize>>,
            end: Vec<usize>,
        }

        impl HLD {
            pub fn new<T: Copy>(tree: Tree<T>, root: usize) -> Self {
                let size = tree.len();
                let mut ret = Self {
                    size,
                    par: vec![None; size],
                    head: vec![0; size],
                    id: vec![0; size],
                    rid: vec![0; size],
                    next: vec![None; size],
                    end: vec![0; size],
                };
                let mut tr = vec![vec![]; size];
                for (i, edges) in tree.nodes.iter().enumerate() {
                    for &TreeEdge { to, .. } in edges.neighbors() {
                        tr[i].push(to);
                    }
                }
                ret.dfs_sub(&mut tr, root, None, &mut vec![1; size]);
                ret.dfs_build(&tr, root, &mut 0);
                ret
            }

            fn dfs_sub(
                &mut self,
                tree: &mut [Vec<usize>],
                cur: usize,
                par: Option<usize>,
                sub: &mut Vec<usize>,
            ) {
                self.par[cur] = par;
                tree[cur].retain(|&x| Some(x) != par);
                let mut t = 0;
                let n = tree[cur].len();
                for i in 0..n {
                    let to = tree[cur][i];
                    self.dfs_sub(tree, to, Some(cur), sub);
                    sub[cur] += sub[to];
                    if sub[to] > t {
                        t = sub[to];
                        self.next[cur] = Some(to);
                        tree[cur].swap(i, 0);
                    }
                }
            }

            fn dfs_build(&mut self, tree: &[Vec<usize>], cur: usize, index: &mut usize) {
                self.id[cur] = *index;
                self.rid[*index] = cur;
                *index += 1;
                for (i, &to) in tree[cur].iter().enumerate() {
                    self.head[to] = if i == 0 { self.head[cur] } else { to };
                    self.dfs_build(tree, to, index);
                }
                self.end[cur] = *index;
            }

            pub fn path_query_vertex(&self, mut x: usize, mut y: usize) -> Vec<(usize, usize)> {
                let mut ret = vec![];
                loop {
                    if self.id[x] > self.id[y] {
                        std::mem::swap(&mut x, &mut y);
                    }
                    ret.push((max(self.id[self.head[y]], self.id[x]), self.id[y] + 1));
                    if self.head[x] == self.head[y] {
                        break;
                    }
                    y = self.par[self.head[y]].unwrap();
                }
                ret
            }

            pub fn path_query_edge(&self, mut x: usize, mut y: usize) -> Vec<(usize, usize)> {
                let mut ret = vec![];
                loop {
                    if self.id[x] > self.id[y] {
                        std::mem::swap(&mut x, &mut y);
                    }
                    if self.head[x] == self.head[y] {
                        if x != y {
                            ret.push((self.id[x] + 1, self.id[y] + 1));
                        }
                        break;
                    }
                    ret.push((self.id[self.head[y]], self.id[y] + 1));
                    y = self.par[self.head[y]].unwrap();
                }
                ret
            }

            pub fn subtree_query_vertex(&self, x: usize) -> (usize, usize) {
                (self.id[x], self.end[x])
            }

            pub fn subtree_query_edge(&self, x: usize) -> (usize, usize) {
                (self.id[x] + 1, self.end[x])
            }

            pub fn parent(&self, x: usize) -> Option<usize> {
                self.par[x]
            }

            pub fn get_id(&self, x: usize) -> usize {
                self.id[x]
            }

            pub fn lca(&self, mut u: usize, mut v: usize) -> usize {
                loop {
                    if self.id[u] > self.id[v] {
                        std::mem::swap(&mut u, &mut v);
                    }
                    if self.head[u] == self.head[v] {
                        return u;
                    }
                    v = self.par[self.head[v]].unwrap();
                }
            }
        }
    }

    #[derive(Clone, Debug)]
    pub struct TreeEdge<T> {
        pub to: usize,
        pub weight: T,
    }

    #[derive(Clone, Debug)]
    pub struct TreeNode<T> {
        pub index: usize,
        pub parent: Option<TreeEdge<T>>,
        pub children: Vec<TreeEdge<T>>,
    }

    #[derive(Clone, Debug)]
    pub struct Tree<T> {
        pub nodes: Vec<TreeNode<T>>,
    }

    impl<T> TreeNode<T> {
        pub fn neighbors(&self) -> impl DoubleEndedIterator<Item = &TreeEdge<T>> {
            self.children.iter().chain(self.parent.iter())
        }

        pub fn neighbors_size(&self) -> usize {
            self.children.len() + self.parent.as_ref().map_or(0, |_| 1)
        }
    }

    impl<T: Copy> Tree<T> {
        pub fn new(size: usize) -> Self {
            Self {
                nodes: (0..size)
                    .map(|i| TreeNode {
                        index: i,
                        parent: None,
                        children: vec![],
                    })
                    .collect(),
            }
        }

        pub fn add_undirected(&mut self, edges: impl IntoIterator<Item = (usize, usize, T)>) {
            for (u, v, w) in edges {
                self.nodes[u].children.push(TreeEdge { to: v, weight: w });
                self.nodes[v].children.push(TreeEdge { to: u, weight: w });
            }
        }

        pub fn add_directed(&mut self, edges: impl IntoIterator<Item = (usize, usize, T)>) {
            for (p, c, w) in edges {
                assert!(self.nodes[c].parent.is_none());
                self.nodes[p].children.push(TreeEdge { to: c, weight: w });
                self.nodes[c].parent = Some(TreeEdge { to: p, weight: w });
            }
        }
    }

    impl<T> Tree<T> {
        pub fn len(&self) -> usize {
            self.nodes.len()
        }

        pub fn is_empty(&self) -> bool {
            self.nodes.is_empty()
        }
    }
}
0