#721133 (Rust) No.899 γatheree

提出ソース
結果

問題	No.899 γatheree
ユーザー	Haar
提出日時	2021-12-01 23:30:37
言語	Rust (1.77.0 + proconio)
結果	RE
実行時間	-
コード長	17,995 bytes
コンパイル時間	13,558 ms
コンパイル使用メモリ	402,656 KB
実行使用メモリ	25,932 KB
最終ジャッジ日時	2024-07-05 01:30:19
合計ジャッジ時間	22,164 ms
ジャッジサーバーID （参考情報）	judge3 / judge5
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 6,812 KB
testcase_01	RE	-
testcase_02	RE	-
testcase_03	RE	-
testcase_04	RE	-
testcase_05	RE	-
testcase_06	RE	-
testcase_07	RE	-
testcase_08	RE	-
testcase_09	RE	-
testcase_10	RE	-
testcase_11	RE	-
testcase_12	RE	-
testcase_13	RE	-
testcase_14	RE	-
testcase_15	RE	-
testcase_16	RE	-
testcase_17	RE	-
testcase_18	RE	-
testcase_19	RE	-
testcase_20	RE	-
testcase_21	AC	958 ms 25,652 KB
testcase_22	AC	966 ms 25,908 KB
testcase_23	AC	1,009 ms 25,932 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
#![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::{
        algebra::update_sum::*,
        ds::lazy_segtree::*,
        //chmin, chmax,
        //mul_vec,
        tree::{depth_query::*, *},
    };

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

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

            input!(cin, n: usize);

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

            let res = TreeDepthQuery::new(&tree, 0);

            input!(cin, a: [u64; n]);
            let mut seg = LazySegmentTree::new(n, UpdateSum::<u64, u64>::new());

            for i in 0..n {
                let (l, r) = res.me_query(i);
                seg.update(l..r, Some(a[i]));
            }

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

                let mut ans = 0;

                let mut f = |(l, r)| {
                    ans += seg.fold(l..r);
                    seg.update(l..r, Some(0));
                };

                if let Some(x) = res.ancestor(x, 2) {
                    f(res.me_query(x));
                }

                if let Some(x) = res.ancestor(x, 1) {
                    f(res.me_query(x));
                    res.children_query(x, 1).map(|p| f(p));
                }

                f(res.me_query(x));

                res.children_query(x, 1).map(|p| f(p));

                res.children_query(x, 2).map(|p| f(p));

                let (l, r) = res.me_query(x);
                seg.update(l..r, Some(ans));

                writeln!(cout, "{}", ans)?;
            }

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

fn main() {
    main::Problem::default().main().unwrap();
}
use crate as haar_lib;
pub mod algebra {
    pub mod action {
        pub trait Action {
            type FType;
            type UType;
            fn fold_id(&self) -> Self::FType;
            fn fold(&self, x: Self::FType, y: Self::FType) -> Self::FType;
            fn update_id(&self) -> Self::UType;
            fn update(&self, next: Self::UType, cur: Self::UType) -> Self::UType;
            fn convert(&self, x: Self::FType, y: Self::UType, l: usize) -> Self::FType;
        }
    }

    pub mod update_sum {
        use crate::algebra::action::Action;

        use std::{
            marker::PhantomData,
            ops::{Add, Mul},
        };

        #[derive(Clone, Default)]
        pub struct UpdateSum<T, U>(PhantomData<T>, PhantomData<U>);

        impl<T, U> UpdateSum<T, U> {
            pub fn new() -> Self {
                Self(PhantomData, PhantomData)
            }
        }

        impl<T, U> Action for UpdateSum<T, U>
        where
            T: Add<Output = T> + Default + From<U>,
            U: Mul<Output = U> + Default + From<u64>,
        {
            type FType = T;

            type UType = Option<U>;

            fn fold_id(&self) -> Self::FType {
                T::default()
            }

            fn fold(&self, x: Self::FType, y: Self::FType) -> Self::FType {
                x + y
            }

            fn update_id(&self) -> Self::UType {
                None
            }

            fn update(&self, x: Self::UType, y: Self::UType) -> Self::UType {
                match x {
                    Some(_) => x,
                    _ => y,
                }
            }

            fn convert(&self, x: Self::FType, y: Self::UType, l: usize) -> Self::FType {
                match y {
                    Some(y) => T::from(y * U::from(l as u64)),
                    _ => x,
                }
            }
        }
    }
}
pub mod algo {
    pub mod bsearch {
        #[doc = " x以上となる最小のindexを求める。"]
        #[doc = ""]
        #[doc = " # Complexity"]
        #[doc = " Time complexity $O(\\log(n))$"]
        pub fn lower_bound<T: Ord>(a: &[T], value: &T) -> usize {
            let n = a.len();
            let mut lb = 0;
            let mut len = n;
            while len > 0 {
                let half = len / 2;
                let mid = lb + half;
                if &a[mid] < value {
                    len -= half + 1;
                    lb = mid + 1;
                } else {
                    len = half;
                }
            }
            lb
        }

        #[doc = " xを超える最小のindexを求める。"]
        #[doc = ""]
        #[doc = " # Complexity"]
        #[doc = " Time complexity $O(\\log(n))$"]
        pub fn upper_bound<T: Ord>(a: &[T], value: &T) -> usize {
            let n = a.len();
            let mut ub = 0;
            let mut len = n;
            while len > 0 {
                let half = len / 2;
                let mid = ub + half;
                if &a[mid] <= value {
                    len -= half + 1;
                    ub = mid + 1;
                } else {
                    len = half;
                }
            }
            ub
        }

        #[doc = " lower_bound, upper_boundの組を求める。"]
        #[doc = ""]
        #[doc = " # Complexity"]
        #[doc = " Time complexity $O(\\log(n))$"]
        pub fn equal_range<T: Ord>(a: &[T], value: &T) -> (usize, usize) {
            (lower_bound(a, value), upper_bound(a, value))
        }
    }
}
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 {
        use crate::algebra::action::Action;

        pub use crate::ds::traits::*;

        use std::ops::Range;

        pub struct LazySegmentTree<T, U, A> {
            size: usize,
            data: Vec<T>,
            lazy: Vec<U>,
            action: A,
        }

        impl<T: Clone + Eq, U: Clone + Eq, A: Clone + Action<FType = T, UType = U>>
            LazySegmentTree<T, U, A>
        {
            pub fn new(n: usize, a: A) -> Self {
                let size = n.next_power_of_two() * 2;
                Self {
                    size,
                    data: vec![a.fold_id(); size],
                    lazy: vec![a.update_id(); size],
                    action: a,
                }
            }

            fn propagate(&mut self, i: usize) {
                if self.lazy[i] == self.action.update_id() {
                    return;
                }
                if i < self.size / 2 {
                    self.lazy[i << 1] = self
                        .action
                        .update(self.lazy[i].clone(), self.lazy[i << 1].clone());
                    self.lazy[i << 1 | 1] = self
                        .action
                        .update(self.lazy[i].clone(), self.lazy[i << 1 | 1].clone());
                }
                let len = (self.size / 2) >> (31 - (i as u32).leading_zeros());
                self.data[i] = self
                    .action
                    .convert(self.data[i].clone(), self.lazy[i].clone(), len);
                self.lazy[i] = self.action.update_id();
            }

            fn propagate_top_down(&mut self, mut i: usize) {
                let mut temp = vec![];
                while i > 1 {
                    i >>= 1;
                    temp.push(i);
                }
                for &i in temp.iter().rev() {
                    self.propagate(i);
                }
            }

            fn bottom_up(&mut self, mut i: usize) {
                while i > 1 {
                    i >>= 1;
                    self.propagate(i << 1);
                    self.propagate(i << 1 | 1);
                    self.data[i] = self
                        .action
                        .fold(self.data[i << 1].clone(), self.data[i << 1 | 1].clone());
                }
            }
        }

        impl<T: Clone + Eq, U: Clone + Eq, A: Clone + Action<FType = T, UType = U>>
            Updatable<Range<usize>> for LazySegmentTree<T, U, A>
        {
            type Value = U;

            fn update(&mut self, Range { start: l, end: r }: Range<usize>, x: U) {
                self.propagate_top_down(l + self.size / 2);
                if r < self.size / 2 {
                    self.propagate_top_down(r + self.size / 2);
                }
                {
                    let mut l = l + self.size / 2;
                    let mut r = r + self.size / 2;
                    while l < r {
                        if r & 1 == 1 {
                            r -= 1;
                            self.lazy[r] = self.action.update(x.clone(), self.lazy[r].clone());
                            self.propagate(r);
                        }
                        if l & 1 == 1 {
                            self.lazy[l] = self.action.update(x.clone(), self.lazy[l].clone());
                            self.propagate(l);
                            l += 1;
                        }
                        r >>= 1;
                        l >>= 1;
                    }
                }
                self.bottom_up(l + self.size / 2);
                if r < self.size / 2 {
                    self.bottom_up(r + self.size / 2);
                }
            }
        }

        impl<T: Clone + Eq, U: Clone + Eq, A: Clone + Action<FType = T, UType = U>>
            FoldableMut<Range<usize>> for LazySegmentTree<T, U, A>
        {
            type Output = T;

            fn fold(&mut self, Range { start: l, end: r }: Range<usize>) -> Self::Output {
                self.propagate_top_down(l + self.size / 2);
                if r < self.size / 2 {
                    self.propagate_top_down(r + self.size / 2);
                }
                let mut ret_l = self.action.fold_id();
                let mut ret_r = self.action.fold_id();
                let mut l = l + self.size / 2;
                let mut r = r + self.size / 2;
                while l < r {
                    if r & 1 == 1 {
                        r -= 1;
                        self.propagate(r);
                        ret_r = self.action.fold(self.data[r].clone(), ret_r.clone());
                    }
                    if l & 1 == 1 {
                        self.propagate(l);
                        ret_l = self.action.fold(ret_l.clone(), self.data[l].clone());
                        l += 1;
                    }
                    r >>= 1;
                    l >>= 1;
                }
                self.action.fold(ret_l, ret_r)
            }
        }
    }
}
pub mod tree {
    pub mod depth_query {
        use crate::{algo::bsearch::lower_bound, tree::*};

        use std::collections::VecDeque;

        pub struct TreeDepthQuery {
            par: Vec<Option<usize>>,
            depth: Vec<usize>,
            left: Vec<usize>,
            right: Vec<usize>,
            bfs_ord: Vec<Vec<usize>>,
            dfs_ord: Vec<Vec<usize>>,
            ord: Vec<usize>,
        }

        impl TreeDepthQuery {
            pub fn new<T>(tree: &Tree<T>, root: usize) -> Self {
                let size = tree.len();
                let mut ret = Self {
                    par: vec![None; size],
                    depth: vec![0; size],
                    left: vec![0; size],
                    right: vec![0; size],
                    bfs_ord: vec![],
                    dfs_ord: vec![],
                    ord: vec![0; size],
                };
                ret.dfs(&tree, root, None, 0, &mut 0);
                let mut q = VecDeque::new();
                q.push_back((root, 0));
                let mut ord = 0;
                while let Some((i, d)) = q.pop_front() {
                    if ret.bfs_ord.len() <= d {
                        ret.bfs_ord.push(vec![]);
                    }
                    ret.bfs_ord[d].push(ord);
                    ret.ord[i] = ord;
                    ord += 1;
                    for &TreeEdge { to, .. } in tree.nodes[i].neighbors() {
                        if Some(to) != ret.par[i] {
                            q.push_back((to, d + 1));
                        }
                    }
                }
                ret
            }

            fn dfs<T>(
                &mut self,
                tree: &Tree<T>,
                cur: usize,
                par: Option<usize>,
                d: usize,
                ord: &mut usize,
            ) {
                self.par[cur] = par;
                self.depth[cur] = d;
                if self.dfs_ord.len() <= d {
                    self.dfs_ord.push(vec![]);
                }
                self.dfs_ord[d].push(*ord);
                self.left[cur] = *ord;
                *ord += 1;
                for &TreeEdge { to, .. } in tree.nodes[cur].neighbors() {
                    if Some(to) != par {
                        self.dfs(tree, to, Some(cur), d + 1, ord);
                    }
                }
                self.right[cur] = *ord;
            }

            pub fn children_query(&self, i: usize, d: usize) -> Option<(usize, usize)> {
                let d = d + self.depth[i];
                if self.bfs_ord.len() > d {
                    let l = lower_bound(&self.dfs_ord[d], &self.left[i]);
                    let r = lower_bound(&self.dfs_ord[d], &self.right[i]);
                    if l >= self.bfs_ord[d].len() {
                        return None;
                    }
                    if r == 1 {
                        return None;
                    }
                    Some((self.bfs_ord[d][l], self.bfs_ord[d][r - 1] + 1))
                } else {
                    None
                }
            }

            pub fn me_query(&self, i: usize) -> (usize, usize) {
                (self.ord[i], self.ord[i] + 1)
            }

            pub fn ancestor(&self, i: usize, k: usize) -> Option<usize> {
                let mut p = i;
                for _ in 0..k {
                    match self.par[p] {
                        Some(x) => p = x,
                        _ => return None,
                    }
                }
                Some(p)
            }
        }
    }

    #[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()
        }
    }
}
yukicoder

結果

テストケース

ソースコード
