#979542 (Rust) No.900 aδδitivee

提出ソース
結果

問題	No.900 aδδitivee
ユーザー	ngtkana
提出日時	2024-05-07 18:41:49
言語	Rust (1.77.0)
結果	AC
実行時間	190 ms / 2,000 ms
コード長	21,846 bytes
コンパイル時間	17,551 ms
コンパイル使用メモリ	378,052 KB
実行使用メモリ	24,576 KB
最終ジャッジ日時	2024-05-07 18:42:17
合計ジャッジ時間	26,028 ms
ジャッジサーバーID （参考情報）	judge1 / judge4
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	0 ms 5,248 KB
testcase_01	AC	1 ms 5,248 KB
testcase_02	AC	1 ms 5,376 KB
testcase_03	AC	1 ms 5,376 KB
testcase_04	AC	0 ms 5,376 KB
testcase_05	AC	1 ms 5,376 KB
testcase_06	AC	1 ms 5,376 KB
testcase_07	AC	148 ms 23,936 KB
testcase_08	AC	160 ms 24,064 KB
testcase_09	AC	164 ms 23,936 KB
testcase_10	AC	178 ms 24,064 KB
testcase_11	AC	152 ms 23,936 KB
testcase_12	AC	151 ms 23,936 KB
testcase_13	AC	151 ms 23,936 KB
testcase_14	AC	149 ms 23,936 KB
testcase_15	AC	148 ms 23,936 KB
testcase_16	AC	156 ms 23,936 KB
testcase_17	AC	190 ms 23,936 KB
testcase_18	AC	181 ms 23,936 KB
testcase_19	AC	147 ms 23,936 KB
testcase_20	AC	150 ms 23,936 KB
testcase_21	AC	153 ms 23,936 KB
testcase_22	AC	153 ms 24,576 KB
testcase_23	AC	145 ms 24,576 KB
testcase_24	AC	148 ms 24,576 KB
testcase_25	AC	138 ms 24,576 KB
testcase_26	AC	160 ms 24,576 KB
testcase_27	AC	139 ms 24,576 KB
testcase_28	AC	132 ms 24,576 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
use crate::segtree::Segtree;
use proconio::input;
use std::cmp::Reverse;
use std::ops::Bound;

fn main() {
    input! {
        n: usize,
        edges: [(usize, usize, i64); n - 1],
        q: usize,
    }
    let mut g = vec![Vec::new(); n];
    for &(i, j, _) in &edges {
        g[i].push(j);
        g[j].push(i);
    }
    let mut stack = vec![0];
    let mut parent = vec![usize::MAX; n];
    let mut depth = vec![0; n];
    parent[0] = 0;
    let mut size = vec![1; n];
    while let Some(x) = stack.pop() {
        if x <= isize::MAX as usize {
            stack.push(!x);
            g[x].retain(|&y| y != parent[x]);
            for &y in &g[x] {
                if parent[y] != usize::MAX {
                    continue;
                }
                parent[y] = x;
                depth[y] = depth[x] + 1;
                stack.push(y);
            }
        } else {
            let x = !x;
            g[x].sort_unstable_by_key(|&y| Reverse(size[y]));
            for &y in &g[x] {
                size[x] += size[y];
            }
        }
    }
    let mut sorted = Vec::new();
    let mut position = vec![usize::MAX; n];
    let mut head = (0..n).collect::<Vec<_>>();
    let mut stack = vec![0];
    while let Some(x) = stack.pop() {
        position[x] = sorted.len();
        sorted.push(x);
        if let Some(&y) = g[x].first() {
            head[y] = head[x];
        }
        stack.extend(g[x].iter().rev().copied());
    }
    let mut original = Segtree::<O1>::from_len(n);
    let mut additional = Segtree::<O2>::from_len(n);
    for &(i, j, w) in &edges {
        let i = if parent[i] == j { i } else { j };
        *original.entry(position[i]) = w;
    }
    for _ in 0..q {
        input! {
            com: u8,
        }
        match com {
            1 => {
                input! {
                    i: usize,
                    x: i64,
                }
                let mut e = additional.entry(position[i]);
                e.additional += x;
                e.weighted += x * depth[i] as i64;
            }
            2 => {
                input! {
                    mut j: usize,
                }
                let mut ans = 0;
                let coeff = depth[j];
                loop {
                    if head[j] == 0 {
                        let original =
                            original.fold((Bound::Excluded(0), Bound::Included(position[j])));
                        let Value {
                            additional,
                            weighted,
                        } = additional.fold(0..=position[j]);
                        ans += additional * coeff as i64 - weighted + original;
                        break;
                    }
                    let original = original.fold(position[head[j]]..=position[j]);
                    let Value {
                        additional,
                        weighted,
                    } = additional.fold(position[head[j]]..=position[j]);
                    ans += additional * coeff as i64 - weighted + original;
                    j = parent[head[j]];
                }
                println!("{}", ans);
            }
            _ => unreachable!(),
        }
    }
}

enum O1 {}
impl segtree::Op for O1 {
    type Value = i64;

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

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

#[derive(Clone, Copy)]
struct Value {
    additional: i64,
    weighted: i64,
}

enum O2 {}
impl segtree::Op for O2 {
    type Value = Value;

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

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

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

#[allow(dead_code)]
mod segtree {
    use core::fmt;
    use std::collections::BTreeMap;
    use std::iter::FromIterator;
    use std::ops::Deref;
    use std::ops::DerefMut;
    use std::ops::Index;
    use std::ops::RangeBounds;
    pub trait Op {
        type Value;
        fn identity() -> Self::Value;
        fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value;
    }
    pub struct Segtree<O: Op> {
        values: Vec<O::Value>,
    }
    impl<O: Op> Segtree<O> {
        pub fn from_len(n: usize) -> Self
        where
            O::Value: Clone,
        {
            Self::new(&vec![O::identity(); n])
        }

        pub fn new(values: &[O::Value]) -> Self
        where
            O::Value: 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 }
        }

        pub fn fold<R: RangeBounds<usize>>(&self, range: R) -> O::Value {
            let n = self.values.len() / 2;
            let (mut start, mut end) = open(range, n);
            assert!(start <= end && end <= n);
            start += n;
            end += n;
            let mut left = O::identity();
            let mut right = O::identity();
            while start < end {
                if start % 2 == 1 {
                    left = O::op(&left, &self.values[start]);
                    start += 1;
                }
                if end % 2 == 1 {
                    end -= 1;
                    right = O::op(&self.values[end], &right);
                }
                start /= 2;
                end /= 2;
            }
            O::op(&left, &right)
        }

        pub fn entry(&mut self, index: usize) -> Entry<O> {
            let n = self.values.len() / 2;
            Entry {
                segtree: self,
                index: n + index,
            }
        }

        pub fn iter(&self) -> impl Iterator<Item = &O::Value> {
            self.values[self.values.len() / 2..].iter()
        }

        pub fn as_slice(&self) -> &[O::Value] {
            &self.values[self.values.len() / 2..]
        }
    }
    impl<O: Op> fmt::Debug for Segtree<O>
    where
        O::Value: fmt::Debug,
    {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("Segtree")
                .field("values", &self.values)
                .finish()
        }
    }
    impl<O: Op> FromIterator<O::Value> for Segtree<O>
    where
        O::Value: Clone,
    {
        fn from_iter<I: IntoIterator<Item = O::Value>>(iter: I) -> Self {
            Self::new(&iter.into_iter().collect::<Vec<_>>())
        }
    }
    impl<O: Op> Index<usize> for Segtree<O> {
        type Output = O::Value;

        fn index(&self, index: usize) -> &Self::Output {
            &self.values[self.values.len() / 2 + index]
        }
    }
    pub struct Entry<'a, O: Op> {
        segtree: &'a mut Segtree<O>,
        index: usize,
    }
    impl<'a, O: Op> Drop for Entry<'a, O> {
        fn drop(&mut self) {
            let mut index = self.index;
            while index != 0 {
                index /= 2;
                self.segtree.values[index] = O::op(
                    &self.segtree.values[index * 2],
                    &self.segtree.values[index * 2 + 1],
                );
            }
        }
    }
    impl<'a, O: Op> Deref for Entry<'a, O> {
        type Target = O::Value;

        fn deref(&self) -> &Self::Target {
            &self.segtree.values[self.index]
        }
    }
    impl<'a, O: Op> DerefMut for Entry<'a, O> {
        fn deref_mut(&mut self) -> &mut Self::Target {
            &mut self.segtree.values[self.index]
        }
    }
    impl<'a, O: Op> fmt::Debug for Entry<'a, O>
    where
        O::Value: fmt::Debug,
    {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("Entry").field("index", &self.index).finish()
        }
    }
    pub struct SparseSegtree<K, O: Op> {
        inner: Segtree<O>,
        keys: Vec<K>,
    }
    impl<K: Ord, O: Op> SparseSegtree<K, O> {
        pub fn new(kv: &[(K, O::Value)]) -> Self
        where
            K: Clone,
            O::Value: Clone,
        {
            let mut keys = kv.iter().map(|(k, _)| k.clone()).collect::<Vec<_>>();
            keys.sort();
            let values = kv.iter().map(|(_, v)| v.clone()).collect::<Vec<_>>();
            Self {
                inner: Segtree::new(&values),
                keys,
            }
        }

        pub fn fold<R: RangeBounds<K>>(&self, range: R) -> O::Value {
            let (start, end) = open_key(range, &self.keys);
            self.inner.fold(start..end)
        }

        pub fn entry(&mut self, key: &K) -> Entry<'_, O> {
            let index = self.keys.binary_search(key).unwrap() + self.keys.len();
            Entry {
                segtree: &mut self.inner,
                index,
            }
        }

        pub fn keys(&self) -> &[K] {
            &self.keys
        }

        pub fn iter(&self) -> impl Iterator<Item = (&K, &O::Value)> {
            self.keys.iter().zip(self.inner.as_slice())
        }

        pub fn collect_map(&self) -> BTreeMap<K, O::Value>
        where
            K: Clone,
            O::Value: Clone,
        {
            self.keys
                .iter()
                .cloned()
                .zip(self.inner.iter().cloned())
                .collect()
        }
    }
    impl<K, O: Op> fmt::Debug for SparseSegtree<K, O>
    where
        K: fmt::Debug,
        O::Value: fmt::Debug,
    {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("SparseSegtree")
                .field("inner", &self.inner)
                .field("keys", &self.keys)
                .finish()
        }
    }
    impl<K: Ord, O: Op> FromIterator<(K, O::Value)> for SparseSegtree<K, O>
    where
        K: Clone,
        O::Value: Clone,
    {
        fn from_iter<I: IntoIterator<Item = (K, O::Value)>>(iter: I) -> Self {
            Self::new(&iter.into_iter().collect::<Vec<_>>())
        }
    }
    impl<K: Ord, O: Op> Index<K> for SparseSegtree<K, O> {
        type Output = O::Value;

        fn index(&self, key: K) -> &Self::Output {
            &self.inner[self.keys.binary_search(&key).unwrap()]
        }
    }
    pub struct Sparse2dSegtree<K, L, O: Op> {
        segtrees: Vec<SparseSegtree<L, O>>,
        keys: Vec<K>,
    }
    impl<K, L, O: Op> Sparse2dSegtree<K, L, O>
    where
        K: Ord + Clone,
        L: Ord + Clone,
        O::Value: Clone,
    {
        pub fn new(points: &[(K, L, O::Value)]) -> Self {
            let mut keys = points.iter().map(|(k, _, _)| k.clone()).collect::<Vec<_>>();
            keys.sort();
            keys.dedup();
            let mut lvs = vec![vec![]; keys.len() * 2];
            for (k, l, v) in points {
                let mut i = keys.binary_search(k).unwrap();
                i += keys.len();
                while i != 0 {
                    lvs[i].push((l.clone(), v.clone()));
                    i /= 2;
                }
            }
            let segtrees = lvs
                .into_iter()
                .map(|lvs_| {
                    let mut ls = lvs_.iter().map(|(l, _)| l).collect::<Vec<_>>();
                    ls.sort();
                    ls.dedup();
                    let mut lvs = ls
                        .iter()
                        .map(|&l| (l.clone(), O::identity()))
                        .collect::<Vec<_>>();
                    for (l, v) in &lvs_ {
                        let i = ls.binary_search(&l).unwrap();
                        lvs[i].1 = O::op(&lvs[i].1, v);
                    }
                    SparseSegtree::new(&lvs)
                })
                .collect::<Vec<_>>();
            Self { segtrees, keys }
        }

        pub fn fold(&self, i: impl RangeBounds<K>, j: impl RangeBounds<L> + Clone) -> O::Value {
            let (mut i0, mut i1) = open_key(i, &self.keys);
            i0 += self.keys.len();
            i1 += self.keys.len();
            let mut left = O::identity();
            let mut right = O::identity();
            while i0 < i1 {
                if i0 % 2 == 1 {
                    left = O::op(&left, &self.segtrees[i0].fold(j.clone()));
                    i0 += 1;
                }
                if i1 % 2 == 1 {
                    i1 -= 1;
                    right = O::op(&self.segtrees[i1].fold(j.clone()), &right);
                }
                i0 /= 2;
                i1 /= 2;
            }
            O::op(&left, &right)
        }

        pub fn apply(&mut self, k: &K, l: &L, mut f: impl FnMut(&mut O::Value)) {
            let mut i = self.keys.binary_search(k).unwrap();
            i += self.keys.len();
            while i != 0 {
                f(&mut self.segtrees[i].entry(l));
                i /= 2;
            }
        }

        pub fn iter(&self) -> impl Iterator<Item = (&K, &L, &O::Value)> {
            self.keys
                .iter()
                .zip(self.segtrees[self.keys.len()..].iter())
                .flat_map(|(k, segtree)| segtree.iter().map(move |(l, v)| (k, l, v)))
        }

        pub fn collect_map(&self) -> BTreeMap<(K, L), O::Value>
        where
            K: Clone,
            L: Clone,
            O::Value: Clone,
        {
            self.keys
                .iter()
                .flat_map(|k| {
                    self.segtrees[self.keys.len() + self.keys.binary_search(k).unwrap()]
                        .iter()
                        .map(move |(l, v)| ((k.clone(), l.clone()), v.clone()))
                })
                .collect()
        }
    }
    impl<K, L, O: Op> fmt::Debug for Sparse2dSegtree<K, L, O>
    where
        K: fmt::Debug,
        L: fmt::Debug,
        O::Value: fmt::Debug,
    {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("Sparse2dSegtree")
                .field("segtrees", &self.segtrees)
                .field("keys", &self.keys)
                .finish()
        }
    }
    impl<K, L, O: Op> FromIterator<(K, L, O::Value)> for Sparse2dSegtree<K, L, O>
    where
        K: Ord + Clone,
        L: Ord + Clone,
        O::Value: Clone,
    {
        fn from_iter<I: IntoIterator<Item = (K, L, O::Value)>>(iter: I) -> Self {
            Self::new(&iter.into_iter().collect::<Vec<_>>())
        }
    }
    impl<K: Ord, L: Ord, O: Op> Index<K> for Sparse2dSegtree<K, L, O> {
        type Output = SparseSegtree<L, O>;

        fn index(&self, i: K) -> &Self::Output {
            &self.segtrees[self.keys.binary_search(&i).unwrap() + self.keys.len()]
        }
    }
    impl<K: Ord, L: Ord, O: Op> Index<(K, L)> for Sparse2dSegtree<K, L, O> {
        type Output = O::Value;

        fn index(&self, (i, j): (K, L)) -> &Self::Output {
            &self.segtrees[self.keys.binary_search(&i).unwrap() + self.keys.len()][j]
        }
    }
    pub struct Dense2dSegtree<O: Op> {
        values: Vec<Vec<O::Value>>,
    }
    impl<O: Op> Dense2dSegtree<O> {
        pub fn new(values: &[Vec<O::Value>]) -> Self
        where
            O::Value: Clone,
        {
            let values_ = values;
            let h = values.len();
            let w = values.get(0).map_or(0, Vec::len);
            let mut values = vec![vec![O::identity(); 2 * w]; 2 * h];
            for (values, values_) in values[h..].iter_mut().zip(values_) {
                values[w..].clone_from_slice(values_);
                for j in (1..w).rev() {
                    values[j] = O::op(&values[j * 2], &values[j * 2 + 1]);
                }
            }
            for i in (1..h).rev() {
                for j in 0..2 * w {
                    values[i][j] = O::op(&values[i * 2][j], &values[i * 2 + 1][j]);
                }
            }
            Self { values }
        }

        pub fn fold(&self, i: impl RangeBounds<usize>, j: impl RangeBounds<usize>) -> O::Value {
            let h = self.values.len() / 2;
            let w = self.values.get(0).map_or(0, |v| v.len() / 2);
            let (mut i0, mut i1) = open(i, h);
            assert!(i0 <= i1 && i1 <= h);
            let (mut j0, mut j1) = open(j, w);
            assert!(j0 <= j1 && j1 <= w);
            i0 += h;
            i1 += h;
            j0 += w;
            j1 += w;
            let mut left = O::identity();
            let mut right = O::identity();
            while i0 < i1 {
                if i0 % 2 == 1 {
                    let mut j0 = j0;
                    let mut j1 = j1;
                    while j0 < j1 {
                        if j0 % 2 == 1 {
                            left = O::op(&left, &self.values[i0][j0]);
                            j0 += 1;
                        }
                        if j1 % 2 == 1 {
                            j1 -= 1;
                            right = O::op(&self.values[i0][j1], &right);
                        }
                        j0 /= 2;
                        j1 /= 2;
                    }
                    i0 += 1;
                }
                if i1 % 2 == 1 {
                    i1 -= 1;
                    let mut j0 = j0;
                    let mut j1 = j1;
                    while j0 < j1 {
                        if j0 % 2 == 1 {
                            left = O::op(&left, &self.values[i1][j0]);
                            j0 += 1;
                        }
                        if j1 % 2 == 1 {
                            j1 -= 1;
                            right = O::op(&self.values[i1][j1], &right);
                        }
                        j0 /= 2;
                        j1 /= 2;
                    }
                }
                i0 /= 2;
                i1 /= 2;
            }
            O::op(&left, &right)
        }

        pub fn entry(&mut self, i: usize, j: usize) -> Dense2dEntry<O> {
            let h = self.values.len() / 2;
            let w = self.values.get(0).map_or(0, |v| v.len() / 2);
            Dense2dEntry {
                segtree: self,
                i: h + i,
                j: w + j,
            }
        }

        pub fn iter(&self) -> impl Iterator<Item = &[O::Value]> {
            self.values[self.values.len() / 2..]
                .iter()
                .map(|v| &v[v.len() / 2..])
        }

        pub fn collect_vec(&self) -> Vec<Vec<O::Value>>
        where
            O::Value: Clone,
        {
            self.iter().map(<[_]>::to_vec).collect()
        }
    }
    impl<O: Op> fmt::Debug for Dense2dSegtree<O>
    where
        O::Value: fmt::Debug,
    {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("Dense2dSegtree")
                .field("values", &self.values)
                .finish()
        }
    }
    impl<O: Op> Index<usize> for Dense2dSegtree<O> {
        type Output = [O::Value];

        fn index(&self, index: usize) -> &Self::Output {
            &self.values[self.values.len() / 2 + index]
        }
    }
    pub struct Dense2dEntry<'a, O: Op> {
        segtree: &'a mut Dense2dSegtree<O>,
        i: usize,
        j: usize,
    }
    impl<'a, O: Op> Drop for Dense2dEntry<'a, O> {
        fn drop(&mut self) {
            let mut i = self.i;
            let mut j = self.j / 2;
            while j != 0 {
                self.segtree.values[i][j] = O::op(
                    &self.segtree.values[i][2 * j],
                    &self.segtree.values[i][2 * j + 1],
                );
                j /= 2;
            }
            i /= 2;
            while i != 0 {
                let mut j = self.j;
                while j != 0 {
                    self.segtree.values[i][j] = O::op(
                        &self.segtree.values[i * 2][j],
                        &self.segtree.values[i * 2 + 1][j],
                    );
                    j /= 2;
                }
                i /= 2;
            }
        }
    }
    impl<'a, O: Op> Deref for Dense2dEntry<'a, O> {
        type Target = O::Value;

        fn deref(&self) -> &Self::Target {
            &self.segtree.values[self.i][self.j]
        }
    }
    impl<'a, O: Op> DerefMut for Dense2dEntry<'a, O> {
        fn deref_mut(&mut self) -> &mut Self::Target {
            &mut self.segtree.values[self.i][self.j]
        }
    }
    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)
    }
    fn open_key<K: Ord, B: RangeBounds<K>>(bounds: B, keys: &[K]) -> (usize, usize) {
        use std::ops::Bound;
        let start = match bounds.start_bound() {
            Bound::Unbounded => 0,
            Bound::Included(x) => match keys.binary_search(x) {
                Ok(i) | Err(i) => i,
            },
            Bound::Excluded(x) => match keys.binary_search(x) {
                Ok(i) => i + 1,
                Err(i) => i,
            },
        };
        let end = match bounds.end_bound() {
            Bound::Unbounded => keys.len(),
            Bound::Included(x) => match keys.binary_search(x) {
                Ok(i) => i + 1,
                Err(i) => i,
            },
            Bound::Excluded(x) => match keys.binary_search(x) {
                Ok(i) | Err(i) => i,
            },
        };
        (start, end)
    }
}
// }}}
yukicoder

結果

テストケース

ソースコード
