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

• 問題: No.900 aδδitivee
• ユーザー: niuez
• 提出日時: 2019-10-20 22:28:40
• 言語: Rust (1.83.0 + proconio)
• 結果: WA
• 実行時間: -
• コード長: 10,115 bytes
• コンパイル時間: 18,455 ms
• コンパイル使用メモリ: 378,152 KB
• 実行使用メモリ: 32,896 KB
• 最終ジャッジ日時: 2024-07-02 17:33:33
• 合計ジャッジ時間: 22,753 ms
• ジャッジサーバーID （参考情報）: judge1 / judge3

ファイルパターン	結果
sample	AC * 2
other	WA * 27

ソースコード

pub struct DirectedTree {
    n: usize,
    r: usize,
    g: Vec<Vec<usize>>,
    p: Vec<Option<(usize, usize)>>,
}

impl DirectedTree {
    pub fn new<I: IntoIterator<Item=(usize, usize)>>(n: usize, es: I) -> Self {
        let mut g: Vec<_> = (0..n).map(|_| Vec::new()).collect();
        let mut p = vec![None; n];
        for (u, v) in es {
            p[v] = Some((u, g[u].len()));
            g[u].push(v);
        }
        DirectedTree {
            n: n,
            r: (0..n).find(|&x| p[x].is_none()).unwrap(),
            g: g,
            p: p,
        }
    }

    pub fn next(&self, v: usize) -> std::slice::Iter<usize> { self.g[v].iter() }
    pub fn parent(&self, v: usize) -> Option<usize> { self.p[v].map(|(u, i)| self.g[u][i]) }
    pub fn root(&self) -> usize { self.r }
    pub fn len(&self) -> usize { self.n }
}

pub struct HeavyLightDecomposition {
    din: Vec<usize>,
    dout: Vec<usize>,
    sz: Vec<usize>,
    seq: Vec<usize>,
    heavy: Vec<Option<usize>>,
    next: Vec<Option<usize>>,
    par: Vec<Option<usize>>,
}

impl HeavyLightDecomposition {
    fn dfs_sz(&mut self, tree: &DirectedTree, v: usize) {
        self.sz[v] = 1;
        for &u in tree.next(v) {
            self.dfs_sz(tree, u);
            self.par[u] = Some(v);
            self.sz[v] += self.sz[u];
        }
        self.heavy[v] = tree.next(v).max_by_key(|&&i| self.sz[i]).map(|&x| x);
    }
    fn dfs_hld(&mut self, tree: &DirectedTree, v: usize, mut t: usize) -> usize {
        self.din[v] = t;
        self.seq.push(v);
        t = t + 1;
        if let Some(h) = self.heavy[v] {
            t = self.dfs_hld(tree, h, t);
            self.next[h] = self.next[v];
            for &u in tree.next(v).filter(|&&u| u != h) {
                t = self.dfs_hld(tree, u, t);
                self.next[u] = Some(u);
            }
        }
        self.dout[v] = t;
        t
    }
    pub fn build(tree: &DirectedTree) -> Self {
        let n = tree.len();
        let mut hld = HeavyLightDecomposition {
            din: vec![0; n],
            dout: vec![0; n],
            seq: Vec::new(),
            sz: vec![0; n],
            heavy: vec![None; n],
            next: vec![None; n],
            par: vec![None; n],
        };
        hld.dfs_sz(tree, tree.root());
        hld.dfs_hld(tree, tree.root(), 0);
        hld
    }
    pub fn sequence(&self) -> std::slice::Iter<usize> { self.seq.iter() }
    pub fn lca(&self, mut v: usize, mut u: usize) -> usize {
        loop {
            if self.din[u] > self.din[v] { std::mem::swap(&mut v, &mut u); }
            if self.next[u] == self.next[v] { break }
            v = self.par[self.next[v].unwrap()].unwrap();
        }
        u
    }
    pub fn path(&self, mut v: usize, mut u: usize, edge: bool) -> (Vec<std::ops::Range<usize>>, Vec<std::ops::Range<usize>>) {
        let mut l = Vec::new();
        let mut r = Vec::new();
        loop {
            if self.din[u] > self.din[v] {
                std::mem::swap(&mut v, &mut u);
                std::mem::swap(&mut l, &mut r);
            }
            if self.next[u] == self.next[v] {
                let e = if edge { 1 } else { 0 };
                l.push(self.din[u] + e..self.din[v] + 1);
                break
            }
            else {
                let ne = self.next[v].unwrap();
                l.push(self.din[ne]..self.din[v] + 1);
                v = self.par[ne].unwrap();
            }
        }
        (l, r)
    }
    pub fn subtree(&self, v: usize, edge: bool) -> std::ops::Range<usize> { self.din[v] + if edge { 1 } else { 0 }..self.dout[v] }
}


pub trait Magma: Sized + Clone {
  fn op(&self, rhs: &Self) -> Self;
}

pub trait Associative: Magma {}

pub trait Unital: Magma {
  fn identity() -> Self;
}

pub trait Monoid: Magma + Associative + Unital {}

pub trait Pow: Magma {
    fn pow(&self, p: usize) -> Self;
}

pub trait Reverse: Magma {
    fn reverse(&self) -> Self;
}

pub trait Inv: Magma {
    fn inv(&self) -> Self;
}

pub trait Effect<E: Monoid> {
    fn effect(&self, e: &E) -> Self;
}

impl<T: Magma + Associative + Unital> Monoid for T {}

use std::ops::Range;

pub struct LazySegmentTree<T: Monoid + Effect<E>, E: Monoid + Pow> {
    node: Vec<T>,
    lazy: Vec<E>,
    sz: usize,
}

impl<T: Monoid + Effect<E>, E: Monoid + Pow> LazySegmentTree<T, E> {
    pub fn init(arr: &[T]) -> Self {
        let mut sz = 1;
        while sz < arr.len() { sz *= 2 }
        let mut node = vec![T::identity(); sz << 1];
        let lazy = vec![E::identity(); sz << 1];
        for i in 0..arr.len() { node[i + sz] = arr[i].clone(); }
        for i in (1..sz).rev() { node[i] = node[i << 1].op(&node[(i << 1) + 1]); }
        Self { node: node, lazy: lazy, sz: sz }
    }

    fn push(&mut self, i: usize, sz: usize) {
        self.node[i] = self.node[i].effect(&self.lazy[i].pow(sz));
        if (i << 1) < self.node.len() {
            self.lazy[i << 1] = self.lazy[i << 1].op(&self.lazy[i]);
            self.lazy[(i << 1) + 1] = self.lazy[(i << 1) + 1].op(&self.lazy[i]);
        }
        self.lazy[i] = E::identity();
    }

    fn update_raw(&mut self, i: usize, a: usize, b: usize, l: usize, r: usize, e: &E) {
        self.push(i, r - l);
        if b <= l || r <= a { return; }
        else if a <= l && r <= b {
            self.lazy[i] = self.lazy[i].op(e);
            self.push(i, r - l);
        }
        else {
            self.update_raw(i << 1, a, b, l, (l + r) >> 1, e);
            self.update_raw((i << 1) + 1, a, b, (l + r) >> 1, r, e);
            self.node[i] = self.node[i << 1].op(&self.node[(i << 1) + 1]);
        }
    }

    pub fn effect(&mut self, ran: Range<usize>, e: E) {
        let sz = self.sz;
        self.update_raw(1, ran.start, ran.end, 0, sz, &e);
    }

    fn fold_raw(&mut self, i: usize, a: usize, b: usize, l: usize, r: usize) -> T {
        self.push(i, r - l);
        if b <= l || r <= a { T::identity() }
        else if a <= l && r <= b { self.node[i].clone() }
        else {
            self.fold_raw(i << 1, a, b, l, (l + r) >> 1)
                .op(&self.fold_raw((i << 1) + 1, a, b, (l + r) >> 1, r))
        }
    }

    pub fn fold(&mut self, ran: Range<usize>) -> T {
        let sz = self.sz;
        self.fold_raw(1, ran.start, ran.end, 0, sz)
    }
}

/// Reading from standard input
///
/// `input!` is useful for competition programming.
/// There are some forms.
///
/// - Tuple
///
/// ```
/// use cp_rust_library::*;
/// input! { source = "2 3", ab: (usize, usize), }
/// assert_eq!(ab.0, 2);
/// assert_eq!(ab.1, 3);
/// ```
///
/// - Array
/// ```
/// use cp_rust_library::*;
/// input! { source = "1 2 3 4", a: [usize; 4], }
/// assert_eq!(a, vec![1, 2, 3, 4]);
/// ```
///
/// - String -> Vec<char>
/// ```
/// use cp_rust_library::*;
/// input! { source = "qwerty", s: chars, }
/// assert_eq!('q', s[0]);
/// ```
///
/// - Other
/// ```
/// use cp_rust_library::*;
/// input! { source = "123", a: usize, }
/// assert_eq!(123, a);
/// ```
/// 
/// This macro will use parse::<$type>() to parse string.

#[macro_export]
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_rec!{ iter, $($r)* }
    };
    ($($r:tt)*) => {
        let s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_rec!{ iter, $($r)* }
    };
}

#[macro_export]
macro_rules! input_rec {
    ($iter:expr) => {};
    ($iter:expr, ) => {};
    ($iter:expr, $var:ident : $t:tt, $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_rec! { $iter, $($r)* }
    };
}

#[macro_export]
macro_rules! read_value {

    // tuple
    ($iter:expr, ( $($t:tt), * )) => {
        ( $(read_value!($iter, $t)), * )
    };
    
    // array
    ($iter:expr, [ $t:tt; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };

    ($iter:expr, query) => {{
        let t = $iter.next().unwrap().parse::<usize>().expect("error");
        if t == 1 {
            (t, read_value!($iter, usize), read_value!($iter, usize))
        }
        else {
            (t, read_value!($iter, usize), 0)
        }
    }};
    
    // string
    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };
    
    // any other
    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };

}


#[derive(Clone)]
struct Am(i64);

impl Magma for Am {
    fn op(&self, right: &Self) -> Self { Am(self.0 + right.0) }
}
impl Associative for Am {}

impl Unital for Am {
    fn identity() -> Self { Am(0) }
}

#[derive(Clone)]
struct Uq(i64);

impl Magma for Uq {
    fn op(&self, right: &Self) -> Self {
        Uq(self.0 + right.0)
    }
}
impl Associative for Uq {}
impl Unital for Uq {
    fn identity() -> Self { Uq(0) }
}
impl Pow for Uq {
    fn pow(&self, len: usize) -> Self { Uq(self.0 * len as i64) }
}
impl Effect<Uq> for Am {
    fn effect(&self, u: &Uq) -> Am {
        Am(self.0 + u.0)
    }
}


fn main() {
    input! {
        n: usize,
        edges: [(usize, usize, usize); n - 1],
        q: usize,
        qs: [query; q],
    }
    let mut w = vec![0; n + 1];
    for &(_, v, a) in edges.iter() {
        w[v] = a as i64;
    }

    let hld = HeavyLightDecomposition::build(
        &DirectedTree::new(n, edges.iter().map(|x| (x.0, x.1)))
        );
    let seq = hld.sequence().map(|&v| w[v]).collect::<Vec<_>>();
    let mut seg = LazySegmentTree::init(
        &seq.into_iter().map(|x| Am(x)).collect::<Vec<_>>()
        );
    for (t, a, b) in qs {
        if t == 1 {
            seg.effect(hld.subtree(a, true), Uq(b as i64));
        }
        else {
            let (l, r) = hld.path(0, a, true);
            println!("{}", 
                    l.iter().fold(Am::identity(), |lm, ran| lm.op(&seg.fold(ran.clone()))).op(
                    &r.iter().fold(Am::identity(), |lm, ran| lm.op(&seg.fold(ran.clone())))).0
                     );
        }
    }
}