#979547 (Rust) No.899 γatheree

• 問題: No.899 γatheree
• ユーザー: ngtkana
• 提出日時: 2024-05-07 19:49:08
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 636 ms / 2,000 ms
• コード長: 9,771 bytes
• コンパイル時間: 15,722 ms
• コンパイル使用メモリ: 376,012 KB
• 実行使用メモリ: 24,740 KB
• 最終ジャッジ日時: 2024-11-30 15:23:37
• 合計ジャッジ時間: 28,164 ms
• ジャッジサーバーID （参考情報）: judge5 / judge4

ファイルパターン	結果
sample	AC * 1
other	AC * 23

ソースコード

use crate::cmpmore::CmpMore;
use crate::lazy_segtree::LazySegtree;
use proconio::input;
use std::collections::VecDeque;

fn main() {
    input! {
        n: usize,
        edges: [(usize, usize); n - 1],
        a: [u64; n],
        q: usize,
        queries: [usize; q],
    }
    let mut g = vec![vec![]; n];
    for (i, j) in edges {
        g[i].push(j);
        g[j].push(i);
    }
    let mut queue = VecDeque::from(vec![0]);
    let mut position = vec![usize::MAX; n];
    let mut sorted = Vec::new();
    let mut parent = vec![usize::MAX; n];
    while let Some(i) = queue.pop_front() {
        position[i] = sorted.len();
        sorted.push(i);
        for &j in &g[i] {
            if position[j] == usize::MAX {
                queue.push_back(j);
                parent[j] = i;
            }
        }
    }
    let mut seg = sorted.iter().map(|&i| a[i]).collect::<LazySegtree<O>>();
    let mut children = vec![(usize::MAX, usize::MIN); n];
    let mut grandchildren = vec![(usize::MAX, usize::MIN); n];
    for &i in sorted.iter().rev() {
        for &j in &g[i] {
            if parent[i] == j {
                continue;
            }
            children[i].0.change_min(position[j]);
            children[i].1.change_max(position[j]);
            grandchildren[i].0.change_min(children[j].0);
            grandchildren[i].1.change_max(children[j].1);
        }
    }
    for &i in &queries {
        let mut ans = 0;
        ans += seg.get(position[i]);
        if children[i].0 != usize::MAX {
            ans += seg.fold(children[i].0..=children[i].1);
            seg.range_apply(children[i].0..=children[i].1, &Some(0));
        }
        if grandchildren[i].0 != usize::MAX {
            ans += seg.fold(grandchildren[i].0..=grandchildren[i].1);
            seg.range_apply(grandchildren[i].0..=grandchildren[i].1, &Some(0));
        }
        let j = parent[i];
        if j != usize::MAX {
            ans += seg.get(position[j]);
            ans -= seg.get(position[i]);
            seg.range_apply(position[j]..=position[j], &Some(0));
            if children[j].0 != usize::MAX {
                ans += seg.fold(children[j].0..=children[j].1);
                seg.range_apply(children[j].0..=children[j].1, &Some(0));
            }
            let k = parent[j];
            if k != usize::MAX {
                ans += seg.get(position[k]);
                seg.range_apply(position[k]..=position[k], &Some(0));
            }
        }
        seg.range_apply(position[i]..=position[i], &Some(ans));
        println!("{}", ans);
    }
}

enum O {}
impl lazy_segtree::Op for O {
    type Operator = Option<u64>;
    type Value = u64;

    fn identity() -> Self::Value {
        0
    }

    fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value {
        lhs + rhs
    }

    fn apply(op: &Self::Operator, value: &Self::Value) -> Self::Value {
        op.unwrap_or(*value)
    }

    fn identity_op() -> Self::Operator {
        None
    }

    fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator {
        op.or(*other)
    }
}

// cmpmore {{{
// https://ngtkana.github.io/ac-adapter-rs/cmpmore/index.html

#[allow(dead_code)]
mod cmpmore {
    pub fn change_min<T: PartialOrd>(lhs: &mut T, rhs: T) {
        if *lhs > rhs {
            *lhs = rhs;
        }
    }
    pub fn change_max<T: PartialOrd>(lhs: &mut T, rhs: T) {
        if *lhs < rhs {
            *lhs = rhs;
        }
    }
    #[macro_export]
    macro_rules! change_min {
        ($lhs:expr, $rhs:expr) => {
            let rhs = $rhs;
            let lhs = $lhs;
            $crate::cmpmore::change_min(lhs, rhs);
        };
    }
    #[macro_export]
    macro_rules! change_max {
        ($lhs:expr, $rhs:expr) => {
            let rhs = $rhs;
            let lhs = $lhs;
            $crate::cmpmore::change_max(lhs, rhs);
        };
    }
    pub trait CmpMore: PartialOrd + Sized {
        fn change_min(&mut self, rhs: Self) {
            change_min(self, rhs)
        }
        fn change_max(&mut self, rhs: Self) {
            change_max(self, rhs)
        }
    }
    impl<T: PartialOrd + Sized> CmpMore for T {}
}
// }}}
// lazy_segtree {{{
// https://ngtkana.github.io/ac-adapter-rs/lazy_segtree/index.html

#[allow(dead_code)]
mod lazy_segtree {
    use std::iter::FromIterator;
    use std::mem::replace;
    use std::ops::RangeBounds;
    pub trait Op {
        type Value;
        type Operator: PartialEq;
        fn identity() -> Self::Value;
        fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value;
        fn apply(op: &Self::Operator, value: &Self::Value) -> Self::Value;
        fn identity_op() -> Self::Operator;
        fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator;
    }
    #[derive(Debug, Clone)]
    pub struct LazySegtree<O: Op> {
        values: Vec<O::Value>,
        operators: Vec<O::Operator>,
    }
    impl<O: Op> LazySegtree<O> {
        pub fn new(values: &[O::Value]) -> Self
        where
            O::Value: Clone,
            O::Operator: Clone,
        {
            let values_ = values;
            let n = values_.len();
            let mut values = vec![O::identity(); 2 * n];
            values[n..].clone_from_slice(values_);
            for i in (1..n).rev() {
                values[i] = O::op(&values[i * 2], &values[i * 2 + 1]);
            }
            Self {
                values,
                operators: vec![O::identity_op(); 2 * n],
            }
        }

        pub fn range_apply<R: RangeBounds<usize>>(&mut self, range: R, f: &O::Operator) {
            let n = self.operators.len() / 2;
            let (l, r) = open(range, n);
            if l == r {
                return;
            }
            let l = n + l;
            let r = n + r;
            for p in (0..usize::BITS - r.leading_zeros()).rev() {
                self.push(l >> p);
                self.push((r - 1) >> p);
            }
            {
                let mut l = l;
                let mut r = r;
                while l < r {
                    if l & 1 != 0 {
                        self.operators[l] = O::compose(f, &self.operators[l]);
                        l += 1;
                    }
                    if r & 1 != 0 {
                        r -= 1;
                        self.operators[r] = O::compose(f, &self.operators[r]);
                    }
                    l >>= 1;
                    r >>= 1;
                }
            }
            for p in 1..usize::BITS - r.leading_zeros() {
                self.update(l >> p);
                self.update((r - 1) >> p);
            }
        }

        pub fn fold<R: RangeBounds<usize>>(&mut self, range: R) -> O::Value {
            let n = self.operators.len() / 2;
            let (mut l, mut r) = open(range, n);
            if l == r {
                return O::identity();
            }
            l += n;
            r += n;
            for p in (0..usize::BITS - r.leading_zeros()).rev() {
                self.push(l >> p);
                self.push((r - 1) >> p);
            }
            for p in 1..usize::BITS - r.leading_zeros() {
                self.update(l >> p);
                self.update((r - 1) >> p);
            }
            let mut left = O::identity();
            let mut right = O::identity();
            while l < r {
                if l & 1 != 0 {
                    left = O::op(&left, &O::apply(&self.operators[l], &self.values[l]));
                    l += 1;
                }
                if r & 1 != 0 {
                    r -= 1;
                    right = O::op(&O::apply(&self.operators[r], &self.values[r]), &right);
                }
                l >>= 1;
                r >>= 1;
            }
            O::op(&left, &right)
        }

        pub fn get(&self, i: usize) -> O::Value
        where
            O::Value: Clone,
        {
            let mut i = self.operators.len() / 2 + i;
            let mut value = self.values[i].clone();
            while i > 0 {
                value = O::apply(&self.operators[i], &value);
                i >>= 1;
            }
            value
        }

        pub fn collect(&self) -> Vec<O::Value>
        where
            O::Value: Clone,
        {
            (0..self.operators.len() / 2).map(|i| self.get(i)).collect()
        }

        fn push(&mut self, i: usize) {
            let a = replace(&mut self.operators[i], O::identity_op());
            self.values[i] = O::apply(&a, &self.values[i]);
            if i < self.operators.len() / 2 {
                self.operators[i << 1] = O::compose(&a, &self.operators[i << 1]);
                self.operators[i << 1 | 1] = O::compose(&a, &self.operators[i << 1 | 1]);
            }
        }

        fn update(&mut self, i: usize) {
            self.values[i] = O::op(
                &O::apply(&self.operators[i << 1], &self.values[i << 1]),
                &O::apply(&self.operators[i << 1 | 1], &self.values[i << 1 | 1]),
            )
        }
    }
    impl<O: Op> FromIterator<O::Value> for LazySegtree<O>
    where
        O::Value: Clone,
        O::Operator: Clone,
    {
        fn from_iter<T: IntoIterator<Item = O::Value>>(iter: T) -> Self {
            Self::new(&iter.into_iter().collect::<Vec<_>>())
        }
    }
    fn open<B: RangeBounds<usize>>(bounds: B, n: usize) -> (usize, usize) {
        use std::ops::Bound;
        let start = match bounds.start_bound() {
            Bound::Unbounded => 0,
            Bound::Included(&x) => x,
            Bound::Excluded(&x) => x + 1,
        };
        let end = match bounds.end_bound() {
            Bound::Unbounded => n,
            Bound::Included(&x) => x + 1,
            Bound::Excluded(&x) => x,
        };
        (start, end)
    }
}
// }}}