#990594 (Rust) No.399 動的な領主

• 問題: No.399 動的な領主
• ユーザー: ngtkana
• 提出日時: 2024-06-24 02:26:20
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 884 ms / 2,000 ms
• コード長: 24,146 bytes
• コンパイル時間: 16,984 ms
• コンパイル使用メモリ: 400,916 KB
• 実行使用メモリ: 30,564 KB
• 最終ジャッジ日時: 2024-06-24 02:26:49
• 合計ジャッジ時間: 23,083 ms
• ジャッジサーバーID （参考情報）: judge1 / judge5

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 6,812 KB
testcase_01	AC	1 ms; 6,812 KB
testcase_02	AC	1 ms; 6,944 KB
testcase_03	AC	1 ms; 6,944 KB
testcase_04	AC	4 ms; 6,944 KB
testcase_05	AC	47 ms; 6,944 KB
testcase_06	AC	838 ms; 21,660 KB
testcase_07	AC	811 ms; 21,532 KB
testcase_08	AC	807 ms; 21,656 KB
testcase_09	AC	879 ms; 21,632 KB
testcase_10	AC	5 ms; 6,944 KB
testcase_11	AC	37 ms; 6,940 KB
testcase_12	AC	609 ms; 22,600 KB
testcase_13	AC	595 ms; 22,716 KB
testcase_14	AC	124 ms; 30,336 KB
testcase_15	AC	156 ms; 30,564 KB
testcase_16	AC	298 ms; 26,092 KB
testcase_17	AC	884 ms; 21,616 KB
testcase_18	AC	839 ms; 21,636 KB

ソースコード

use lazy_segtree::LazySegtree;
use proconio::input;
use proconio::marker::Usize1;

fn main() {
    input! {
        n: usize,
        edges: [(Usize1, Usize1); n - 1],
        q: usize,
        queries: [(Usize1, Usize1); q],
    }

    let (hld, _g) = Hld::from_edges(0, &edges);

    let mut lazy_segtree = LazySegtree::<O>::from_iter(vec![Value { len: 1, sum: 1 }; n]);
    let mut ans = 0;
    for &(i, j) in &queries {
        hld.visit_path_segments_including_lca_by_index(i, j, |i, j| {
            ans += lazy_segtree.fold(i..=j).sum;
            lazy_segtree.range_apply(i..=j, &1);
        });
    }
    println!("{ans}");
}

#[derive(Clone, Copy, Debug, PartialEq, Default)]
struct Value {
    len: usize,
    sum: usize,
}

enum O {}
impl lazy_segtree::Op for O {
    type Operator = usize;
    type Value = Value;

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

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

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

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

    fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator {
        op + other
    }
}

pub struct Hld {
    pub parent: Vec<usize>,
    pub index: Vec<usize>,
    pub head: Vec<usize>,
}
impl Hld {
    pub fn from_short_parents(mut parent: Vec<usize>) -> (Self, Vec<Vec<usize>>) {
        parent.insert(0, 0);
        let mut g = vec![Vec::new(); parent.len()];
        for (i, &p) in parent.iter().enumerate().skip(1) {
            g[p].push(i);
        }
        (__build_hld(0, &mut g, parent), g)
    }

    pub fn from_edges(root: usize, edges: &[(usize, usize)]) -> (Self, Vec<Vec<usize>>) {
        let mut g = vec![Vec::new(); edges.len() + 1];
        for &(i, j) in edges {
            g[i].push(j);
            g[j].push(i);
        }
        let parent = __remove_parent(root, &mut g);
        (__build_hld(root, &mut g, parent), g)
    }

    pub fn visit_path_segments_including_lca(
        &self,
        mut i: usize,                    // id
        mut j: usize,                    // id
        mut f: impl FnMut(usize, usize), // id
    ) {
        while self.head[i] != self.head[j] {
            if self.index[i] < self.index[j] {
                f(self.head[j], j);
                j = self.parent[self.head[j]];
            } else {
                f(self.head[i], i);
                i = self.parent[self.head[i]];
            }
        }
        if self.index[i] < self.index[j] {
            f(i, j)
        } else {
            f(j, i)
        }
    }

    pub fn visit_path_segments_including_lca_by_index(
        &self,
        i: usize,                        // id
        j: usize,                        // id
        mut f: impl FnMut(usize, usize), // index
    ) {
        self.visit_path_segments_including_lca(i, j, |i, j| f(self.index[i], self.index[j]));
    }

    pub fn lca(&self, mut i: usize, mut j: usize) -> usize {
        while self.head[i] != self.head[j] {
            if self.index[i] < self.index[j] {
                j = self.parent[self.head[j]];
            } else {
                i = self.parent[self.head[i]];
            }
        }
        std::cmp::min_by_key(i, j, |&i| self.index[i])
    }
}

fn __build_hld(root: usize, g: &mut [Vec<usize>], parent: Vec<usize>) -> Hld {
    let n = g.len();
    __heavy_first(0, g);
    let mut index = vec![usize::MAX; n];
    let mut head = vec![usize::MAX; n];
    head[root] = root;
    __head_and_index(0, &*g, &mut head, &mut index, &mut (0..));
    Hld {
        parent,
        index,
        head,
    }
}
fn __head_and_index(
    i: usize,
    g: &[Vec<usize>],
    head: &mut [usize],
    index: &mut [usize],
    current: &mut std::ops::RangeFrom<usize>,
) {
    index[i] = current.next().unwrap();
    for &j in &g[i] {
        head[j] = if j == g[i][0] { head[i] } else { j };
        __head_and_index(j, g, head, index, current);
    }
}

fn __heavy_first(i: usize, g: &mut [Vec<usize>]) -> usize {
    let mut max = 0;
    let mut size = 1;
    for e in 0..g[i].len() {
        let csize = __heavy_first(g[i][e], g);
        if max < csize {
            max = csize;
            g[i].swap(0, e);
        }
        size += csize;
    }
    size
}

fn __remove_parent(root: usize, g: &mut [Vec<usize>]) -> Vec<usize> {
    let mut stack = vec![root];
    let mut parent = vec![usize::MAX; g.len()];
    parent[root] = root;
    while let Some(i) = stack.pop() {
        g[i].retain(|&j| parent[i] != j);
        for &j in &g[i] {
            parent[j] = i;
            stack.push(j);
        }
    }
    parent
}
// link_cut_tree {{{
// https://ngtkana.github.io/ac-adapter-rs/link_cut_tree/index.html

#[allow(dead_code)]
mod link_cut_tree {
    mod base {
        #[doc(hidden)]
        pub trait OpBase {
            type Value: Clone;
            type InternalValue: Clone;
            fn identity() -> Self::InternalValue;
            fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue;
            fn into_front(value: Self::InternalValue) -> Self::Value;
            fn from_front(value: Self::Value) -> Self::InternalValue;
            fn rev(value: &mut Self::InternalValue);
        }
        pub struct LinkCutTreeBase<O: OpBase> {
            nodes: Vec<Node<O>>,
        }
        impl<O: OpBase> LinkCutTreeBase<O> {
            pub fn new(n: usize) -> Self {
                Self {
                    nodes: (0..n)
                        .map(|id| Node {
                            id,
                            parent: std::ptr::null_mut(),
                            left: std::ptr::null_mut(),
                            right: std::ptr::null_mut(),
                            rev: false,
                            value: O::identity(),
                            acc: O::identity(),
                        })
                        .collect(),
                }
            }

            pub fn from_values(values: impl IntoIterator<Item = O::Value>) -> Self {
                Self {
                    nodes: values
                        .into_iter()
                        .map(O::from_front)
                        .enumerate()
                        .map(|(id, value)| Node {
                            id,
                            parent: std::ptr::null_mut(),
                            left: std::ptr::null_mut(),
                            right: std::ptr::null_mut(),
                            rev: false,
                            value: value.clone(),
                            acc: value,
                        })
                        .collect(),
                }
            }

            pub fn link(&mut self, p: usize, c: usize) {
                unsafe {
                    let c = std::ptr::addr_of_mut!(self.nodes[c]);
                    let p = std::ptr::addr_of_mut!(self.nodes[p]);
                    expose(c);
                    assert!((*c).left.is_null(), "c = {} is not a root", (*c).id);
                    expose(p);
                    assert!(
                        (*c).parent.is_null(),
                        "c = {} and p = {} are already connected",
                        (*c).id,
                        (*p).id
                    );
                    (*c).parent = p;
                    (*p).right = c;
                    update(p);
                }
            }

            pub fn undirected_link(&mut self, i: usize, j: usize) -> bool {
                if self.undirected_is_connected(i, j) {
                    return false;
                }
                self.evert(j);
                self.link(i, j);
                true
            }

            pub fn cut(&mut self, x: usize) -> Option<usize> {
                unsafe {
                    let x = std::ptr::addr_of_mut!(self.nodes[x]);
                    expose(x);
                    let p = (*x).left;
                    (*x).left = std::ptr::null_mut();
                    let ans = p.as_ref().map(|p| p.id);
                    if !p.is_null() {
                        (*p).parent = std::ptr::null_mut();
                    }
                    update(x);
                    ans
                }
            }

            pub fn undirected_cut(&mut self, i: usize, j: usize) -> bool {
                if !self.undirected_has_edge(i, j) {
                    return false;
                }
                self.evert(i);
                self.cut(j);
                true
            }

            pub fn evert(&mut self, x: usize) {
                unsafe {
                    let x = std::ptr::addr_of_mut!(self.nodes[x]);
                    expose(x);
                    rev(x);
                    push(x);
                }
            }

            pub fn undirected_has_edge(&mut self, x: usize, y: usize) -> bool {
                self.parent(x) == Some(y) || self.parent(y) == Some(x)
            }

            pub fn undirected_is_connected(&mut self, x: usize, y: usize) -> bool {
                if x == y {
                    return true;
                }
                unsafe {
                    let x = std::ptr::addr_of_mut!(self.nodes[x]);
                    let y = std::ptr::addr_of_mut!(self.nodes[y]);
                    expose(x);
                    expose(y);
                    !(*x).parent.is_null()
                }
            }

            pub fn lca(&mut self, x: usize, y: usize) -> Option<usize> {
                if x == y {
                    return Some(x);
                }
                unsafe {
                    let x = std::ptr::addr_of_mut!(self.nodes[x]);
                    let y = std::ptr::addr_of_mut!(self.nodes[y]);
                    expose(x);
                    let lca = expose(y);
                    if (*x).parent.is_null() {
                        None
                    } else {
                        Some((*lca).id)
                    }
                }
            }

            pub fn set(&mut self, x: usize, mut f: impl FnMut(O::Value) -> O::Value) {
                unsafe {
                    let x = std::ptr::addr_of_mut!(self.nodes[x]);
                    expose(x);
                    (*x).value = O::from_front(f(O::into_front((*x).value.clone())));
                    update(x);
                }
            }

            pub fn fold(&mut self, x: usize) -> O::Value {
                unsafe {
                    let x = std::ptr::addr_of_mut!(self.nodes[x]);
                    expose(x);
                    O::into_front((*x).acc.clone())
                }
            }

            pub fn undirected_fold(&mut self, i: usize, j: usize) -> Option<O::Value> {
                if !self.undirected_is_connected(i, j) {
                    return None;
                }
                self.evert(i);
                Some(self.fold(j))
            }

            pub fn parent(&mut self, x: usize) -> Option<usize> {
                unsafe {
                    let x = std::ptr::addr_of_mut!(self.nodes[x]);
                    expose(x);
                    let mut p = (*x).left.as_mut()?;
                    while let Some(next) = p.right.as_mut() {
                        p = next;
                    }
                    splay(p);
                    Some(p.id)
                }
            }
        }
        #[derive(Clone, Copy)]
        struct Node<O: OpBase> {
            id: usize,
            parent: *mut Self,
            left: *mut Self,
            right: *mut Self,
            rev: bool,
            value: O::InternalValue,
            acc: O::InternalValue,
        }
        unsafe fn is_splay_root<O: OpBase>(x: *mut Node<O>) -> bool {
            let x = &*x;
            let p = match x.parent.as_ref() {
                Some(p) => p,
                None => return true,
            };
            !std::ptr::eq(x, p.left) && !std::ptr::eq(x, p.right)
        }
        unsafe fn push<O: OpBase>(x: *mut Node<O>) {
            let x = &mut *x;
            if x.rev {
                if let Some(l) = x.left.as_mut() {
                    rev(l);
                }
                if let Some(r) = x.right.as_mut() {
                    rev(r);
                }
                x.rev = false;
            }
        }
        unsafe fn update<O: OpBase>(x: *mut Node<O>) {
            let x = &mut *x;
            x.acc = x.value.clone();
            if !x.left.is_null() {
                x.acc = O::mul(&(*x.left).acc, &x.acc);
            }
            if !x.right.is_null() {
                x.acc = O::mul(&x.acc, &(*x.right).acc);
            }
        }
        unsafe fn rev<O: OpBase>(x: *mut Node<O>) {
            let x = &mut *x;
            std::mem::swap(&mut x.left, &mut x.right);
            O::rev(&mut x.acc);
            x.rev ^= true;
        }
        unsafe fn expose<O: OpBase>(x: *mut Node<O>) -> *mut Node<O> {
            let mut last = std::ptr::null_mut();
            let mut current = x;
            while !current.is_null() {
                splay(current);
                (*current).right = last;
                update(current);
                last = current;
                current = (*current).parent;
            }
            splay(x);
            last
        }
        unsafe fn splay<O: OpBase>(x: *mut Node<O>) {
            let x = &mut *x;
            push(x);
            while !is_splay_root(x) {
                let p = &mut *x.parent;
                if is_splay_root(p) {
                    push(p);
                    push(x);
                    if std::ptr::eq(p.left, x) {
                        rotate_right(p);
                    } else {
                        rotate_left(p);
                    }
                } else {
                    let g = &mut *p.parent;
                    push(g);
                    push(p);
                    push(x);
                    #[allow(clippy::collapsible_else_if)]
                    if std::ptr::eq(p.left, x) {
                        if std::ptr::eq(g.left, p) {
                            rotate_right(g);
                            rotate_right(p);
                        } else {
                            rotate_right(p);
                            rotate_left(g);
                        }
                    } else {
                        if std::ptr::eq(g.left, p) {
                            rotate_left(p);
                            rotate_right(g);
                        } else {
                            rotate_left(g);
                            rotate_left(p);
                        }
                    }
                }
            }
        }
        unsafe fn rotate_left<O: OpBase>(l: *mut Node<O>) {
            let l = &mut *l;
            let r = &mut *l.right;
            let p = l.parent;
            let c = r.left;
            l.right = c;
            if !c.is_null() {
                (*c).parent = l;
            }
            r.left = l;
            l.parent = r;
            r.parent = p;
            update(l);
            update(r);
            if !p.is_null() {
                if std::ptr::eq((*p).left, l) {
                    (*p).left = r;
                } else if std::ptr::eq((*p).right, l) {
                    (*p).right = r;
                }
                update(&mut *p);
            }
        }
        unsafe fn rotate_right<O: OpBase>(r: *mut Node<O>) {
            let r = &mut *r;
            let l = &mut *r.left;
            let p = r.parent;
            let c = l.right;
            r.left = c;
            if !c.is_null() {
                (*c).parent = r;
            }
            l.right = r;
            r.parent = l;
            l.parent = p;
            update(r);
            update(l);
            if !p.is_null() {
                if std::ptr::eq((*p).left, r) {
                    (*p).left = l;
                } else if std::ptr::eq((*p).right, r) {
                    (*p).right = l;
                }
                update(&mut *p);
            }
        }
    }
    pub use base::LinkCutTreeBase;
    use base::OpBase;
    pub trait Op {
        type Value: Clone;
        fn identity() -> Self::Value;
        fn mul(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value;
    }
    impl OpBase for () {
        type InternalValue = ();
        type Value = ();

        fn identity() -> Self::InternalValue {}

        fn mul(_lhs: &Self::InternalValue, _rhs: &Self::InternalValue) -> Self::InternalValue {}

        fn rev(_value: &mut Self::InternalValue) {}

        fn into_front(_value: Self::InternalValue) {}

        fn from_front(_value: Self::Value) -> Self::InternalValue {}
    }
    pub type LinkCutTree = LinkCutTreeBase<()>;
    pub type CommutLinkCutTree<T> = LinkCutTreeBase<Commut<T>>;
    #[doc(hidden)]
    pub struct Commut<T: Op>(T);
    impl<T: Op> OpBase for Commut<T> {
        type InternalValue = T::Value;
        type Value = T::Value;

        fn identity() -> Self::InternalValue {
            T::identity()
        }

        fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue {
            T::mul(lhs, rhs)
        }

        fn rev(_value: &mut Self::InternalValue) {}

        fn into_front(value: Self::InternalValue) -> Self::Value {
            value
        }

        fn from_front(value: Self::Value) -> Self::InternalValue {
            value
        }
    }
    #[doc(hidden)]
    pub struct NonCommut<T: Op>(T);
    pub type NonCommutLinkCutTree<T> = LinkCutTreeBase<NonCommut<T>>;
    impl<T: Op> OpBase for NonCommut<T> {
        type InternalValue = (T::Value, T::Value);
        type Value = T::Value;

        fn identity() -> Self::InternalValue {
            (T::identity(), T::identity())
        }

        fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue {
            (T::mul(&lhs.0, &rhs.0), T::mul(&rhs.1, &lhs.1))
        }

        fn rev(value: &mut Self::InternalValue) {
            std::mem::swap(&mut value.0, &mut value.1);
        }

        fn into_front(value: Self::InternalValue) -> Self::Value {
            value.0
        }

        fn from_front(value: Self::Value) -> Self::InternalValue {
            (value.clone(), value)
        }
    }
}
// }}}
// 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)
    }
}
// }}}