結果

問題 No.899 γatheree
ユーザー HaarHaar
提出日時 2022-09-06 16:54:59
言語 Rust
(1.77.0)
結果
AC  
実行時間 1,658 ms / 2,000 ms
コード長 19,113 bytes
コンパイル時間 19,777 ms
コンパイル使用メモリ 379,480 KB
実行使用メモリ 29,312 KB
最終ジャッジ日時 2024-05-01 13:21:18
合計ジャッジ時間 40,774 ms
ジャッジサーバーID
(参考情報)
judge2 / judge4
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
5,248 KB
testcase_01 AC 2 ms
5,248 KB
testcase_02 AC 2 ms
5,248 KB
testcase_03 AC 2 ms
5,248 KB
testcase_04 AC 16 ms
5,376 KB
testcase_05 AC 2 ms
5,376 KB
testcase_06 AC 1,541 ms
29,168 KB
testcase_07 AC 1,615 ms
29,132 KB
testcase_08 AC 1,300 ms
29,012 KB
testcase_09 AC 1,550 ms
29,276 KB
testcase_10 AC 1,567 ms
29,140 KB
testcase_11 AC 1,522 ms
29,160 KB
testcase_12 AC 1,633 ms
29,112 KB
testcase_13 AC 1,505 ms
29,144 KB
testcase_14 AC 1,548 ms
29,184 KB
testcase_15 AC 1,658 ms
29,088 KB
testcase_16 AC 1,570 ms
29,144 KB
testcase_17 AC 1,543 ms
29,312 KB
testcase_18 AC 1,505 ms
29,128 KB
testcase_19 AC 1,482 ms
29,188 KB
testcase_20 AC 1,395 ms
29,208 KB
testcase_21 AC 983 ms
29,168 KB
testcase_22 AC 988 ms
29,296 KB
testcase_23 AC 982 ms
29,300 KB
権限があれば一括ダウンロードができます

ソースコード

diff #

// Bundled at 2022/09/06 16:54:35 +09:00
// Author: Haar

pub mod main {
    use super::*;
    use haar_lib::{
        algebra::update_sum::*,
        ds::lazy_segtree::*,
        tree::{depth_query::*, *},
        utils::fastio::*,
    };
    #[derive(Clone, Default)]
    pub struct Problem {}
    impl Problem {
        pub fn main(&mut self) -> Result<(), Box<dyn std::error::Error>> {
            let mut io = FastIO::new();
            let n = io.read_u64() as usize;
            let mut tree = Tree::new(n);
            for _ in 0..n - 1 {
                let u = io.read_u64() as usize;
                let v = io.read_u64() as usize;
                tree.extend(Some(TreeEdge::new(u, v, (), ())));
            }
            let res = TreeDepthQuery::new(&tree, 0);
            let a = (0..n).map(|_| io.read_u64()).collect::<Vec<_>>();
            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]));
            }
            let q = io.read_u64() as usize;
            for _ in 0..q {
                let x = io.read_u64() as 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).into_iter().for_each(|p| f(p));
                }
                f(res.me_query(x));
                res.children_query(x, 1).into_iter().for_each(|p| f(p));
                res.children_query(x, 2).into_iter().for_each(|p| f(p));
                let (l, r) = res.me_query(x);
                seg.update(l..r, Some(ans));
                io.writeln(ans);
            }
            Ok(())
        }
    }
}
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 ds {
    pub mod traits {
        pub trait Foldable<Idx> {
            type Output;
            fn fold(&self, range: Idx) -> Self::Output;
        }
        pub trait Foldable2D<Idx> {
            type Output;
            fn fold(&self, x_range: Idx, y_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 Indexable<Idx> {
            type Output;
            fn get(&self, i: Idx) -> Self::Output;
        }
    }
    pub mod lazy_segtree {
        use crate::algebra::action::Action;
        pub use crate::ds::traits::Updatable;
        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());
                }
            }

            pub fn fold(&mut self, Range { start: l, end: r }: Range<usize>) -> T {
                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)
            }
        }
        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);
                }
            }
        }
    }
}
pub mod tree {
    pub mod depth_query {
        use crate::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<E: TreeEdgeTrait>(tree: &Tree<E>, 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 e in tree.nodes[i].neighbors() {
                        if Some(e.to()) != ret.par[i] {
                            q.push_back((e.to(), d + 1));
                        }
                    }
                }
                ret
            }

            fn dfs<E: TreeEdgeTrait>(
                &mut self,
                tree: &Tree<E>,
                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 e in tree.nodes[cur].neighbors() {
                    if Some(e.to()) != par {
                        self.dfs(tree, e.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 = match self.dfs_ord[d].binary_search(&self.left[i]) {
                        Ok(x) | Err(x) => x,
                    };
                    let r = match self.dfs_ord[d].binary_search(&self.right[i]) {
                        Ok(x) | Err(x) => x,
                    };
                    if l >= self.bfs_ord[d].len() {
                        return None;
                    }
                    if r == 0 {
                        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)
            }
        }
    }
    pub trait TreeEdgeTrait {
        type Weight;
        fn from(&self) -> usize;
        fn to(&self) -> usize;
        fn weight(&self) -> Self::Weight;
        fn rev(self) -> Self;
    }
    #[derive(Clone, Debug)]
    pub struct TreeEdge<T, I> {
        pub from: usize,
        pub to: usize,
        pub weight: T,
        pub index: I,
    }
    impl<T, I> TreeEdge<T, I> {
        pub fn new(from: usize, to: usize, weight: T, index: I) -> Self {
            Self {
                from,
                to,
                weight,
                index,
            }
        }
    }
    impl<T: Clone, I> TreeEdgeTrait for TreeEdge<T, I> {
        type Weight = T;

        #[inline]

        fn from(&self) -> usize {
            self.from
        }

        #[inline]

        fn to(&self) -> usize {
            self.to
        }

        #[inline]

        fn weight(&self) -> Self::Weight {
            self.weight.clone()
        }

        fn rev(mut self) -> Self {
            std::mem::swap(&mut self.from, &mut self.to);
            self
        }
    }
    #[derive(Clone, Debug)]
    pub struct TreeNode<E> {
        pub parent: Option<E>,
        pub children: Vec<E>,
    }
    #[derive(Clone, Debug)]
    pub struct Tree<E> {
        pub nodes: Vec<TreeNode<E>>,
    }
    impl<E: TreeEdgeTrait> TreeNode<E> {
        pub fn neighbors(&self) -> impl DoubleEndedIterator<Item = &E> {
            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<E: TreeEdgeTrait + Clone> Tree<E> {
        pub fn new(size: usize) -> Self {
            Self {
                nodes: vec![
                    TreeNode {
                        parent: None,
                        children: vec![],
                    };
                    size
                ],
            }
        }

        pub fn extend(&mut self, edges: impl IntoIterator<Item = E>) {
            for e in edges {
                self.nodes[e.from()].children.push(e.clone());
                self.nodes[e.to()].children.push(e.rev());
            }
        }

        pub fn extend_rooted(&mut self, edges: impl IntoIterator<Item = E>) {
            for e in edges {
                assert!(self.nodes[e.to()].parent.is_none());
                self.nodes[e.from()].children.push(e.clone());
                self.nodes[e.to()].parent.replace(e.rev());
            }
        }
    }
    impl<E> Tree<E> {
        pub fn len(&self) -> usize {
            self.nodes.len()
        }

        pub fn is_empty(&self) -> bool {
            self.nodes.is_empty()
        }
    }
}
pub mod utils {
    pub mod fastio {
        use std::fmt::Display;
        use std::io::{Read, Write};
        pub struct FastIO {
            in_bytes: Vec<u8>,
            in_cur: usize,
            out_buf: std::io::BufWriter<std::io::Stdout>,
        }
        impl FastIO {
            pub fn new() -> Self {
                let mut s = vec![];
                std::io::stdin().read_to_end(&mut s).unwrap();
                let cout = std::io::stdout();
                Self {
                    in_bytes: s,
                    in_cur: 0,
                    out_buf: std::io::BufWriter::new(cout),
                }
            }

            #[inline]
            pub fn getc(&mut self) -> Option<u8> {
                if self.in_cur < self.in_bytes.len() {
                    self.in_cur += 1;
                    Some(self.in_bytes[self.in_cur])
                } else {
                    None
                }
            }

            #[inline]
            pub fn peek(&self) -> Option<u8> {
                if self.in_cur < self.in_bytes.len() {
                    Some(self.in_bytes[self.in_cur])
                } else {
                    None
                }
            }

            #[inline]
            pub fn skip(&mut self) {
                while self.peek().map_or(false, |c| c.is_ascii_whitespace()) {
                    self.in_cur += 1;
                }
            }

            pub fn read_u64(&mut self) -> u64 {
                self.skip();
                let mut ret: u64 = 0;
                while self.peek().map_or(false, |c| c.is_ascii_digit()) {
                    ret = ret * 10 + (self.in_bytes[self.in_cur] - b'0') as u64;
                    self.in_cur += 1;
                }
                ret
            }

            pub fn read_i64(&mut self) -> i64 {
                self.skip();
                let mut ret: i64 = 0;
                let minus = if self.peek() == Some(b'-') {
                    self.in_cur += 1;
                    true
                } else {
                    false
                };
                while self.peek().map_or(false, |c| c.is_ascii_digit()) {
                    ret = ret * 10 + (self.in_bytes[self.in_cur] - b'0') as i64;
                    self.in_cur += 1;
                }
                if minus {
                    ret = -ret;
                }
                ret
            }

            pub fn read_chars(&mut self) -> Vec<char> {
                self.skip();
                let mut ret = vec![];
                while self.peek().map_or(false, |c| c.is_ascii_graphic()) {
                    ret.push(self.in_bytes[self.in_cur] as char);
                    self.in_cur += 1;
                }
                ret
            }

            pub fn write<T: Display>(&mut self, s: T) {
                self.out_buf.write_all(format!("{}", s).as_bytes()).unwrap();
            }

            pub fn writeln<T: Display>(&mut self, s: T) {
                self.write(s);
                self.out_buf.write_all(&[b'\n']).unwrap();
            }
        }
        impl Drop for FastIO {
            fn drop(&mut self) {
                self.out_buf.flush().unwrap();
            }
        }
    }
}
0