#823088 (Rust) No.2163 LCA Sum Query

• 問題: No.2163 LCA Sum Query
• ユーザー: akakimidori
• 提出日時: 2022-12-07 04:36:36
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 935 ms / 6,000 ms
• コード長: 18,906 bytes
• コンパイル時間: 4,301 ms
• コンパイル使用メモリ: 185,080 KB
• 実行使用メモリ: 21,704 KB
• 最終ジャッジ日時: 2024-04-24 12:36:34
• 合計ジャッジ時間: 18,939 ms
• ジャッジサーバーID （参考情報）: judge3 / judge4

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 5,248 KB
testcase_01	AC	1 ms; 5,376 KB
testcase_02	AC	1 ms; 5,376 KB
testcase_03	AC	1 ms; 5,376 KB
testcase_04	AC	2 ms; 5,376 KB
testcase_05	AC	2 ms; 5,376 KB
testcase_06	AC	1 ms; 5,376 KB
testcase_07	AC	1 ms; 5,376 KB
testcase_08	AC	1 ms; 5,376 KB
testcase_09	AC	2 ms; 5,376 KB
testcase_10	AC	2 ms; 5,376 KB
testcase_11	AC	2 ms; 5,376 KB
testcase_12	AC	306 ms; 7,424 KB
testcase_13	AC	394 ms; 17,344 KB
testcase_14	AC	298 ms; 18,448 KB
testcase_15	AC	115 ms; 5,376 KB
testcase_16	AC	242 ms; 16,260 KB
testcase_17	AC	322 ms; 10,352 KB
testcase_18	AC	207 ms; 5,760 KB
testcase_19	AC	141 ms; 5,376 KB
testcase_20	AC	20 ms; 8,576 KB
testcase_21	AC	122 ms; 9,312 KB
testcase_22	AC	688 ms; 21,228 KB
testcase_23	AC	556 ms; 21,400 KB
testcase_24	AC	647 ms; 21,272 KB
testcase_25	AC	543 ms; 21,152 KB
testcase_26	AC	935 ms; 21,196 KB
testcase_27	AC	710 ms; 21,192 KB
testcase_28	AC	926 ms; 21,192 KB
testcase_29	AC	707 ms; 21,192 KB
testcase_30	AC	262 ms; 21,164 KB
testcase_31	AC	265 ms; 21,036 KB
testcase_32	AC	298 ms; 21,572 KB
testcase_33	AC	269 ms; 21,296 KB
testcase_34	AC	307 ms; 21,512 KB
testcase_35	AC	248 ms; 21,568 KB
testcase_36	AC	327 ms; 21,564 KB
testcase_37	AC	272 ms; 21,520 KB
testcase_38	AC	352 ms; 21,704 KB
testcase_39	AC	306 ms; 21,160 KB
testcase_40	AC	357 ms; 21,572 KB
testcase_41	AC	294 ms; 21,160 KB

ソースコード

use std::io::*;

fn main() {
    input! {
        n: usize,
        q: usize,
        edge: [(usize, usize); n - 1],
        query: [(usize, usize, usize); q],
    }
    let n = n + 1;
    let mut hld = HLD::new(n);
    hld.add_edge(0, 1);
    for &(a, b) in edge.iter() {
        hld.add_edge(a, b);
    }
    hld.build(0);
    let mut cnt = SegmentTreePURQ::new(n, 0, |a, b| *a + *b);
    let mut fix_root = LazySegmentTree::new(n, FixRoot);
    let mut move_root = LazySegmentTree::new(n, MoveRoot);
    for i in 1..n {
        let x = hld.vertex(i);
        let parent = hld.parent(x).unwrap();
        move_root.get_mut(i, |p| p.3 = x as i64 - parent as i64);
    }
    let mut state = vec![0i64; n];
    let mut range = vec![];
    let out = std::io::stdout();
    let mut out = std::io::BufWriter::new(out.lock());
    for &(u, r, v) in query.iter() {
        let sign = (state[u] ^ 1) - state[u];
        move_root.update(0, n, (sign, 0));
        let mut sub = state[u];
        hld.path_root(u, &mut range);
        for &(l, r) in range.iter().rev() {
            move_root.update(l, r, (-sign, sign));
            let v = hld.vertex(r - 1);
            let p = hld.sub_tree(v);
            let c = cnt.find(p.0, p.1) - sub;
            fix_root.get_mut(r - 1, |p| {
                let v = v as i64;
                p.0 += sign * c * v;
                p.1 += sign * v;
            });
            fix_root.update(l, r - 1, sign);
            let p = hld.sub_tree(hld.vertex(l));
            sub = cnt.find(p.0, p.1);
        }
        state[u] ^= 1;
        cnt.update(hld.sub_tree(u).0, state[u]);
        let a = hld.sub_tree(r);
        let b = hld.sub_tree(v);
        let ans = if b.0 <= a.0 && a.1 <= b.1 {
            let mut all = fix_root.find(0, n).0;
            hld.path_root(v, &mut range);
            for &(l, r) in range.iter() {
                all += move_root.find(l, r).0;
            }
            if r != v {
                let x = hld.next(v, r);
                let p = hld.sub_tree(x);
                let c = cnt.find(0, n);
                let d = cnt.find(p.0, p.1);
                all -= v as i64 * (c - d) * d;
                all -= fix_root.find(p.0, p.1).0;
            }
            all
        } else {
            fix_root.find(b.0, b.1).0
        };
        writeln!(out, "{}", ans).ok();
    }
}

struct FixRoot;
impl TE for FixRoot {
    type T = (i64, i64);
    type E = i64;
    fn fold(&self, l: &Self::T, r: &Self::T) -> Self::T {
        (l.0 + r.0, l.1 + r.1)
    }
    fn eval(&self, x: &Self::T, f: &Self::E) -> Self::T {
        (x.0 + x.1 * *f, x.1)
    }
    fn merge(&self, g: &Self::E, h: &Self::E) -> Self::E {
        *g + *h
    }
    fn e(&self) -> Self::T {
        (0, 0)
    }
    fn id(&self) -> Self::E {
        0
    }
}

struct MoveRoot;
impl TE for MoveRoot {
    // sum cab, sum ca, sum cb, sum c
    type T = (i64, i64, i64, i64);
    // a, b
    type E = (i64, i64);
    fn fold(&self, l: &Self::T, r: &Self::T) -> Self::T {
        (l.0 + r.0, l.1 + r.1, l.2 + r.2, l.3 + r.3)
    }
    fn eval(&self, x: &Self::T, f: &Self::E) -> Self::T {
        (
            x.0 + f.0 * x.2 + f.1 * x.1 + f.0 * f.1 * x.3,
            x.1 + f.0 * x.3,
            x.2 + f.1 * x.3,
            x.3,
        )
    }
    fn merge(&self, g: &Self::E, h: &Self::E) -> Self::E {
        (g.0 + h.0, g.1 + h.1)
    }
    fn e(&self) -> Self::T {
        (0, 0, 0, 0)
    }
    fn id(&self) -> Self::E {
        (0, 0)
    }
}

// ---------- begin segment tree Point Update Range Query ----------
pub struct SegmentTreePURQ<T, F> {
    n: usize,
    size: usize,
    data: Vec<T>,
    e: T,
    op: F,
}

impl<T, F> SegmentTreePURQ<T, F>
where
    T: Clone,
    F: Fn(&T, &T) -> T,
{
    pub fn new(n: usize, e: T, op: F) -> Self {
        assert!(n > 0);
        let size = n.next_power_of_two();
        let data = vec![e.clone(); 2 * size];
        SegmentTreePURQ {
            n,
            size,
            data,
            e,
            op,
        }
    }
    pub fn update_tmp(&mut self, x: usize, v: T) {
        assert!(x < self.n);
        self.data[x + self.size] = v;
    }
    pub fn update_all(&mut self) {
        for i in (1..self.size).rev() {
            self.data[i] = (self.op)(&self.data[2 * i], &self.data[2 * i + 1]);
        }
    }
    pub fn update(&mut self, x: usize, v: T) {
        assert!(x < self.n);
        let mut x = x + self.size;
        self.data[x] = v;
        x >>= 1;
        while x > 0 {
            self.data[x] = (self.op)(&self.data[2 * x], &self.data[2 * x + 1]);
            x >>= 1;
        }
    }
    pub fn find(&self, l: usize, r: usize) -> T {
        assert!(l <= r && r <= self.n);
        if l == r {
            return self.e.clone();
        }
        let mut l = self.size + l;
        let mut r = self.size + r;
        let mut x = self.e.clone();
        let mut y = self.e.clone();
        while l < r {
            if l & 1 == 1 {
                x = (self.op)(&x, &self.data[l]);
                l += 1;
            }
            if r & 1 == 1 {
                r -= 1;
                y = (self.op)(&self.data[r], &y);
            }
            l >>= 1;
            r >>= 1;
        }
        (self.op)(&x, &y)
    }
    pub fn max_right<P>(&self, l: usize, f: P) -> usize
    where
        P: Fn(&T) -> bool,
    {
        assert!(l <= self.n);
        assert!(f(&self.e));
        if l == self.n {
            return self.n;
        }
        let mut l = l + self.size;
        let mut sum = self.e.clone();
        while {
            l >>= l.trailing_zeros();
            let v = (self.op)(&sum, &self.data[l]);
            if !f(&v) {
                while l < self.size {
                    l <<= 1;
                    let v = (self.op)(&sum, &self.data[l]);
                    if f(&v) {
                        sum = v;
                        l += 1;
                    }
                }
                return l - self.size;
            }
            sum = v;
            l += 1;
            l.count_ones() > 1
        } {}
        self.n
    }
    pub fn min_left<P>(&self, r: usize, f: P) -> usize
    where
        P: Fn(&T) -> bool,
    {
        assert!(r <= self.n);
        assert!(f(&self.e));
        if r == 0 {
            return 0;
        }
        let mut r = r + self.size;
        let mut sum = self.e.clone();
        while {
            r -= 1;
            while r > 1 && r & 1 == 1 {
                r >>= 1;
            }
            let v = (self.op)(&self.data[r], &sum);
            if !f(&v) {
                while r < self.size {
                    r = 2 * r + 1;
                    let v = (self.op)(&self.data[r], &sum);
                    if f(&v) {
                        sum = v;
                        r -= 1;
                    }
                }
                return r + 1 - self.size;
            }
            sum = v;
            (r & (!r + 1)) != r
        } {}
        0
    }
}
// ---------- end segment tree Point Update Range Query ----------
// ---------- begin Heavy-Light decomposition ----------
pub struct HLD {
    size: usize,
    edge: Vec<(usize, usize)>,
    child: Vec<Vec<usize>>,
    path_root: Vec<usize>,
    parent: Vec<usize>,
    left: Vec<usize>,
    right: Vec<usize>,
    inverse: Vec<usize>,
}

impl HLD {
    pub fn new(size: usize) -> Self {
        assert!(size <= 10usize.pow(8));
        HLD {
            size: size,
            edge: Vec::with_capacity(size - 1),
            child: Vec::new(),
            path_root: Vec::new(),
            parent: Vec::new(),
            left: Vec::new(),
            right: Vec::new(),
            inverse: Vec::new(),
        }
    }
    pub fn add_edge(&mut self, a: usize, b: usize) {
        assert!(a != b && a < self.size && b < self.size);
        self.edge.push((a, b));
    }
    pub fn build(&mut self, root: usize) {
        assert!(self.edge.len() + 1 == self.size);
        let size = self.size;
        let mut cnt = vec![0; size];
        for &(a, b) in self.edge.iter() {
            cnt[a] += 1;
            cnt[b] += 1;
        }
        let mut child = cnt
            .into_iter()
            .map(|c| Vec::with_capacity(c))
            .collect::<Vec<_>>();
        for &(a, b) in self.edge.iter() {
            child[a].push(b);
            child[b].push(a);
        }
        let mut parent = vec![size; size];
        let mut q = Vec::with_capacity(size);
        q.push(root);
        parent[root] = root;
        for i in 0..size {
            let v = q[i];
            for u in child[v].clone() {
                assert!(parent[u] == size);
                parent[u] = v;
                child[u].retain(|e| *e != v);
                q.push(u);
            }
        }
        let mut sum = vec![1; size];
        for &v in q.iter().rev() {
            let child = &mut child[v];
            if !child.is_empty() {
                let (pos, _) = child.iter().enumerate().max_by_key(|p| sum[*p.1]).unwrap();
                child.swap(0, pos);
                sum[v] = 1 + child.iter().fold(0, |s, a| s + sum[*a]);
            }
        }
        let mut path_root = (0..size).collect::<Vec<_>>();
        let mut left = vec![0; size];
        let mut right = vec![0; size];
        let mut dfs = vec![(root, false)];
        let mut id = 0;
        while let Some((v, end)) = dfs.pop() {
            if end {
                right[v] = id;
                continue;
            }
            left[v] = id;
            id += 1;
            dfs.push((v, true));
            let child = &child[v];
            if !child.is_empty() {
                for &u in child[1..].iter().rev() {
                    path_root[u] = u;
                    dfs.push((u, false));
                }
                let u = child[0];
                path_root[u] = path_root[v];
                dfs.push((u, false));
            }
        }
        let mut inverse = vec![size; size];
        for (i, l) in left.iter().enumerate() {
            inverse[*l] = i;
        }
        self.child = child;
        self.parent = parent;
        self.left = left;
        self.right = right;
        self.path_root = path_root;
        self.inverse = inverse;
    }
    pub fn lca(&self, mut a: usize, mut b: usize) -> usize {
        assert!(a < self.size && b < self.size);
        let path = &self.path_root;
        let parent = &self.parent;
        let index = &self.left;
        while path[a] != path[b] {
            if index[a] > index[b] {
                std::mem::swap(&mut a, &mut b);
            }
            b = parent[path[b]];
        }
        std::cmp::min((index[a], a), (index[b], b)).1
    }
    pub fn path_root(&self, mut v: usize, res: &mut Vec<(usize, usize)>) {
        assert!(v < self.size);
        res.clear();
        let parent = &self.parent;
        let path_root = &self.path_root;
        let index = &self.left;
        while index[path_root[v]] != 0 {
            let p = path_root[v];
            res.push((index[p], index[v] + 1));
            v = parent[p];
        }
        res.push((0, index[v] + 1));
        res.reverse();
    }
    pub fn path(
        &self,
        src: usize,
        dst: usize,
        up: &mut Vec<(usize, usize)>,
        down: &mut Vec<(usize, usize)>,
    ) {
        assert!(src < self.size && dst < self.size);
        up.clear();
        down.clear();
        let path = &self.path_root;
        let parent = &self.parent;
        let index = &self.left;
        let mut x = src;
        let mut y = dst;
        while path[x] != path[y] {
            if index[x] > index[y] {
                let p = path[x];
                assert!(p == path[p]);
                up.push((index[p], index[x] + 1));
                x = parent[p];
            } else {
                let p = path[y];
                assert!(p == path[p]);
                down.push((index[p], index[y] + 1));
                y = parent[p];
            }
        }
        if index[x] >= index[y] {
            up.push((index[y], index[x] + 1));
        } else {
            down.push((index[x], index[y] + 1));
        }
        down.reverse();
    }
    pub fn sub_tree(&self, v: usize) -> (usize, usize) {
        assert!(v < self.size);
        (self.left[v], self.right[v])
    }
    pub fn parent(&self, v: usize) -> Option<usize> {
        assert!(v < self.size);
        let p = self.parent[v];
        if p == v {
            None
        } else {
            Some(p)
        }
    }
    // s -> t へのパスの2番目の頂点を返す
    pub fn next(&self, s: usize, t: usize) -> usize {
        assert!(s < self.size && t < self.size && s != t);
        let (a, b) = self.sub_tree(s);
        let (c, d) = self.sub_tree(t);
        if !(a <= c && d <= b) {
            return self.parent[s];
        }
        let mut pos = t;
        let mut pre = t;
        while self.path_root[s] != self.path_root[pos] {
            pre = self.path_root[pos];
            pos = self.parent[pre];
        }
        if s == pos {
            pre
        } else {
            self.child[s][0]
        }
    }
    pub fn vertex(&self, x: usize) -> usize {
        assert!(x < self.size);
        self.inverse[x]
    }
}
// ---------- end Heavy-Light decomposition ----------
// ---------- begin Lazy Segment Tree ----------
pub trait TE {
    type T: Clone;
    type E: Clone;
    fn fold(&self, l: &Self::T, r: &Self::T) -> Self::T;
    fn eval(&self, x: &Self::T, f: &Self::E) -> Self::T;
    fn merge(&self, g: &Self::E, h: &Self::E) -> Self::E;
    fn e(&self) -> Self::T;
    fn id(&self) -> Self::E;
}

pub struct LazySegmentTree<R: TE> {
    n: usize,
    size: usize,
    bit: u32,
    op: R,
    data: Vec<(R::T, R::E)>,
}

impl<R: TE> LazySegmentTree<R> {
    pub fn new(n: usize, op: R) -> Self {
        assert!(n > 0);
        let size = n.next_power_of_two();
        let bit = size.trailing_zeros();
        let data = vec![(op.e(), op.id()); 2 * size];
        Self {
            n,
            size,
            bit,
            op,
            data,
        }
    }
    pub fn build<I>(init: I, n: usize, op: R) -> Self
    where
        I: Iterator<Item = R::T>,
    {
        let mut seg = Self::new(n, op);
        for (data, ini) in seg.data[seg.size..].iter_mut().zip(init) {
            data.0 = ini;
        }
        for i in (1..seg.size).rev() {
            seg.pull(i);
        }
        seg
    }
    pub fn update(&mut self, l: usize, r: usize, f: R::E) {
        assert!(l <= r && r <= self.n);
        if l == r {
            return;
        }
        self.push_range(l, r);
        let mut s = l + self.size;
        let mut t = r + self.size;
        while s < t {
            if s & 1 == 1 {
                self.apply(s, &f);
                s += 1;
            }
            if t & 1 == 1 {
                t -= 1;
                self.apply(t, &f);
            }
            s >>= 1;
            t >>= 1;
        }
        let l = l + self.size;
        let r = r + self.size;
        for k in 1..=self.bit {
            if (l >> k) << k != l {
                self.pull(l >> k);
            }
            if (r >> k) << k != r {
                self.pull((r - 1) >> k);
            }
        }
    }
    pub fn find(&mut self, l: usize, r: usize) -> R::T {
        assert!(l <= r && r <= self.n);
        if l == r {
            return self.op.e();
        }
        self.push_range(l, r);
        let mut l = l + self.size;
        let mut r = r + self.size;
        let mut p = self.op.e();
        let mut q = self.op.e();
        while l < r {
            if l & 1 == 1 {
                p = self.op.fold(&p, &self.data[l].0);
                l += 1;
            }
            if r & 1 == 1 {
                r -= 1;
                q = self.op.fold(&self.data[r].0, &q);
            }
            l >>= 1;
            r >>= 1;
        }
        self.op.fold(&p, &q)
    }
    pub fn get_mut<F>(&mut self, x: usize, f: F)
    where
        F: FnOnce(&mut R::T),
    {
        assert!(x < self.n);
        let x = x + self.size;
        for k in (1..=self.bit).rev() {
            self.push(x >> k);
        }
        f(&mut self.data[x].0);
        for k in 1..=self.bit {
            self.pull(x >> k);
        }
    }
    pub fn set_at(&mut self, x: usize, v: R::T) {
        assert!(x < self.n);
        let x = x + self.size;
        for k in (1..=self.bit).rev() {
            self.push(x >> k);
        }
        self.data[x].0 = v;
        for k in 1..=self.bit {
            self.pull(x >> k);
        }
    }
    fn push_range(&mut self, l: usize, r: usize) {
        let l = l + self.size;
        let r = r + self.size;
        for k in (1..=self.bit).rev() {
            if (l >> k) << k != l {
                self.push(l >> k);
            }
            if (r >> k) << k != r {
                self.push((r - 1) >> k);
            }
        }
    }
    fn apply(&mut self, x: usize, f: &R::E) {
        self.data[x].0 = self.op.eval(&self.data[x].0, f);
        self.data[x].1 = self.op.merge(&self.data[x].1, f);
    }
    fn push(&mut self, x: usize) {
        let f = std::mem::replace(&mut self.data[x].1, self.op.id());
        self.apply(2 * x, &f);
        self.apply(2 * x + 1, &f);
    }
    fn pull(&mut self, x: usize) {
        self.data[x].0 = self.op.fold(&self.data[2 * x].0, &self.data[2 * x + 1].0);
    }
}
// ---------- end Lazy Segment Tree ----------
// ---------- begin input macro ----------
// reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
#[macro_export]
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{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_inner!{iter, $($r)*}
    };
}

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

#[macro_export]
macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };
    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };
    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };
    ($iter:expr, bytes) => {
        read_value!($iter, String).bytes().collect::<Vec<u8>>()
    };
    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };
    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}
// ---------- end input macro ----------