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

• 問題: No.399 動的な領主
• ユーザー: ngtkana
• 提出日時: 2021-08-05 00:47:21
• 言語: Rust (1.77.0)
• 結果: WA
• 実行時間: -
• コード長: 20,893 bytes
• コンパイル時間: 1,459 ms
• コンパイル使用メモリ: 161,236 KB
• 実行使用メモリ: 24,576 KB
• 最終ジャッジ日時: 2023-10-14 21:27:11
• 合計ジャッジ時間: 5,186 ms
• ジャッジサーバーID （参考情報）: judge15 / judge14

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 4,352 KB
testcase_01	AC	1 ms; 4,352 KB
testcase_02	WA	-
testcase_03	WA	-
testcase_04	WA	-
testcase_05	WA	-
testcase_06	WA	-
testcase_07	WA	-
testcase_08	WA	-
testcase_09	WA	-
testcase_10	WA	-
testcase_11	WA	-
testcase_12	WA	-
testcase_13	WA	-
testcase_14	AC	77 ms; 24,576 KB
testcase_15	WA	-
testcase_16	WA	-
testcase_17	WA	-
testcase_18	WA	-

ソースコード

#[allow(unused_imports)]
#[cfg(feature = "dbg")]
use dbg::lg;
use {
    make_graph::tuple_make_undirected,
    sort_tree::sort_tree,
    std::{
        iter::repeat_with,
        ptr::{self, null_mut},
    },
};

fn main() {
    let mut buf = ngtio::with_stdin();
    let n = buf.usize();
    let g = tuple_make_undirected(
        n,
        &repeat_with(|| (buf.usize() - 1, buf.usize() - 1))
            .take(n - 1)
            .collect::<Vec<_>>(),
    );
    let [_, parent] = sort_tree(0, &g);
    let mut lct = LinkCutTree::new(n);
    for (i, &p) in parent.iter().enumerate() {
        if i != p {
            lct.link(i, p);
        }
    }
    for _ in 0..buf.usize() {
        let u = buf.usize() - 1;
        let v = buf.usize() - 1;
        let q = lct.lca(u, v).unwrap();
        lct.act_from_root(u, 1);
        lct.act_from_root(v, 1);
        lct.act_from_root(q, -1);
        if let Some(q) = lct.parent(q) {
            lct.act_from_root(q, -1);
        }
    }
    let ans = (0..n)
        .map(|i| lct.value(i) as i64)
        .map(|x| x * (x + 1) / 2)
        .sum::<i64>();
    println!("{}", ans);
}

pub struct LinkCutTree {
    nodes: Vec<*mut Node>,
}
impl LinkCutTree {
    pub fn new(len: usize) -> Self {
        let nodes = (0..len)
            .map(|id| Box::into_raw(Box::new(Node::new(id))))
            .collect::<Vec<_>>();
        Self { nodes }
    }
    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty()
    }
    pub fn len(&self) -> usize {
        self.nodes.len()
    }
    /// 頂点 `child` が根でないとき親の ID を返し、そうでないとき `None` を
    /// 返します。
    ///
    /// # Panics
    ///
    /// - `child` が `self.len()` 以上のとき
    ///
    pub fn parent(&self, child: usize) -> Option<usize> {
        let child = self.__index(child);
        unsafe {
            expose(child);
            let l = (*child).left;
            l.as_ref().map(|l| l.id)
        }
    }
    /// 頂点 `u`, `v` が同じ木に属するとき、`true`
    ///
    /// # Panics
    ///
    /// - `u`, `v` のいずれかが `self.len()` 以上のとき
    ///
    pub fn same(&self, u: usize, v: usize) -> bool {
        let u = self.__index(u);
        let v = self.__index(v);
        unsafe {
            expose(u);
            expose(v);
            ptr::eq(u, v) || !(*u).parent.is_null()
        }
    }
    /// ある木の根 `child` を別の木の頂点 `parent` の子にします。
    ///
    /// # Panics
    ///
    /// - `child`, `parent` のいずれかが `self.len()` 以上のとき
    /// - `child` と `parent` が同じ木に属するとき
    /// - `child` が根でないとき
    ///
    pub fn link(&mut self, child: usize, parent: usize) {
        debug_assert!(!self.same(child, parent));
        let child = self.__index(child);
        let parent = self.__index(parent);
        unsafe { (*child).parent = parent }
    }
    /// 根でない頂点 `child` と、その親を結ぶ辺を消去し、もともと
    /// 親だった頂点の ID を返します。
    ///
    /// # Panics
    ///
    /// - `child` が `self.len()` 以上のとき
    /// - `child` がある頂点の値であるとき
    ///
    pub fn cut(&mut self, child: usize) -> usize {
        debug_assert!(child < self.nodes.len());
        let child = self.nodes[child];
        unsafe {
            expose(child);
            let l = (*child).left;
            debug_assert!(!l.is_null());
            (*l).parent = null_mut();
            (*child).left = null_mut();
            (*child).update();
            (&*l).id
        }
    }
    /// 頂点 `u` と `v` が同じ木に属するとき LCA を返し、そうでないとき
    /// `None` を返します。
    ///
    /// # Panics
    ///
    /// - `u`, `v` のいずれかが `self.len()` 以上のとき
    ///
    pub fn lca(&self, u: usize, v: usize) -> Option<usize> {
        let u = self.__index(u);
        let v = self.__index(v);
        unsafe {
            expose(u);
            expose(v).as_ref().map(|q| q.id)
        }
    }
    pub fn depth(&self, _u: usize) -> usize {
        todo!()
    }
    pub fn act_from_root(&mut self, u: usize, lazy: i32) {
        let u = self.__index(u);
        unsafe {
            expose(u);
            (*u).lazy += lazy;
            (*u).push();
        }
    }
    pub fn value(&self, u: usize) -> i32 {
        let u = self.__index(u);
        unsafe {
            expose(u);
            (&*u).value
        }
    }
    fn __index(&self, index: usize) -> *mut Node {
        if self.nodes.len() <= index {
            panic!(
                "index {} out of range for link-cut tree of length {}",
                index,
                self.len(),
            );
        }
        self.nodes[index]
    }
}

unsafe fn expose(v: *mut Node) -> *mut Node {
    let mut last = null_mut();
    let mut now = v;
    while !now.is_null() {
        (*now).splay();
        (*now).right = last;
        last = now;
        now = (*now).parent;
    }
    (*v).splay();
    last
}

struct Node {
    left: *mut Node,
    right: *mut Node,
    parent: *mut Node,
    id: usize,
    lazy: i32,
    value: i32,
}

impl Node {
    fn new(id: usize) -> Self {
        Self {
            left: null_mut(),
            right: null_mut(),
            parent: null_mut(),
            id,
            value: 0,
            lazy: 0,
        }
    }
    unsafe fn rotate(&mut self) {
        let p = self.parent;
        let pp = (*p).parent;
        let c;
        if ptr::eq((*p).left, self) {
            c = self.right;
            self.right = p;
            (*p).left = c;
        } else {
            c = self.left;
            self.left = p;
            (*p).right = c;
        }
        if !pp.is_null() && ptr::eq((*pp).left, p) {
            (*pp).left = self;
        }
        if !pp.is_null() && ptr::eq((*pp).right, p) {
            (*pp).right = self;
        }
        self.parent = pp;
        (*p).parent = self;
        if !c.is_null() {
            (*c).parent = p;
        }
        (*p).update();
        self.update();
    }
    unsafe fn splay(&mut self) {
        while self.state() != 0 {
            let p = self.parent;
            if (*p).state() == 0 {
                (*p).push();
                self.push();
                self.rotate();
            } else {
                let g = (*p).parent;
                (*g).push();
                (*p).push();
                self.push();
                if self.state() == (*p).state() {
                    (*p).rotate();
                    self.rotate();
                } else {
                    self.rotate();
                    self.rotate();
                }
            }
        }
        self.push();
    }
    unsafe fn state(&self) -> i8 {
        if self.parent.is_null() {
            0
        } else if ptr::eq((*self.parent).left, self) {
            1
        } else if ptr::eq((*self.parent).right, self) {
            -1
        } else {
            0
        }
    }
    unsafe fn update(&mut self) {}
    unsafe fn push(&mut self) {
        if self.lazy != 0 {
            if let Some(l) = self.left.as_mut() {
                l.lazy += self.lazy;
            }
            if let Some(r) = self.right.as_mut() {
                r.lazy += self.lazy;
            }
            self.value += self.lazy;
            self.lazy = 0;
        }
    }
}
// sort_tree {{{
#[allow(dead_code)]
mod sort_tree {
    pub fn remove_parent(g: &mut [Vec<usize>], parent: &[usize]) {
        g.iter_mut().enumerate().for_each(|(x, gx)| {
            if let Some(i) = gx.iter().position(|&y| y == parent[x]) {
                gx.swap_remove(i);
            };
        });
    }
    pub fn sort_tree_remove_parent(root: usize, g: &mut [Vec<usize>]) -> [Vec<usize>; 2] {
        let [ord, parent] = sort_tree(root, g);
        remove_parent(g, &parent);
        [ord, parent]
    }
    pub fn sort_tree(root: usize, g: &[Vec<usize>]) -> [Vec<usize>; 2] {
        sort_tree_by(root, g, |x| *x)
    }
    pub fn sort_tree_by<E>(root: usize, g: &[Vec<E>], to: impl Fn(&E) -> usize) -> [Vec<usize>; 2] {
        let mut ord = Vec::new();
        let mut parent = vec![root; g.len()];
        sort_tree_impl(root, root, g, &to, &mut parent, &mut ord);
        [ord, parent]
    }
    fn sort_tree_impl<E>(
        x: usize,
        p: usize,
        g: &[Vec<E>],
        to: &impl Fn(&E) -> usize,
        parent: &mut [usize],
        ord: &mut Vec<usize>,
    ) {
        ord.push(x);
        parent[x] = p;
        g[x].iter()
            .map(to)
            .filter(|&y| y != p)
            .for_each(|y| sort_tree_impl(y, x, g, to, parent, ord))
    }
}
// }}}
// make_graph {{{
#[allow(dead_code)]
mod make_graph {
    pub fn tuple_make_undirected(n: usize, edges: &[(usize, usize)]) -> Vec<Vec<usize>> {
        make_undirected_by(n, edges, |&(u, v)| [u, v])
    }
    pub fn array_make_undirected(n: usize, edges: &[[usize; 2]]) -> Vec<Vec<usize>> {
        make_undirected_by(n, edges, |&[u, v]| [u, v])
    }
    pub fn make_undirected_by<E>(
        n: usize,
        edges: &[E],
        f: impl Fn(&E) -> [usize; 2],
    ) -> Vec<Vec<usize>> {
        let mut g = vec![Vec::new(); n];
        for [u, v] in edges.iter().map(f) {
            g[u].push(v);
            g[v].push(u);
        }
        g
    }
    pub fn tuple_make_directed(n: usize, edges: &[(usize, usize)]) -> Vec<Vec<usize>> {
        make_directed_by(n, edges, |&(u, v)| [u, v])
    }
    pub fn array_make_directed(n: usize, edges: &[[usize; 2]]) -> Vec<Vec<usize>> {
        make_directed_by(n, edges, |&[u, v]| [u, v])
    }
    pub fn make_directed_by<E>(
        n: usize,
        edges: &[E],
        f: impl Fn(&E) -> [usize; 2],
    ) -> Vec<Vec<usize>> {
        let mut g = vec![Vec::new(); n];
        edges.iter().map(f).for_each(|[u, v]| g[u].push(v));
        g
    }
    pub fn tuple_make_undirected_weighted<T: Copy>(
        n: usize,
        edges: &[(usize, usize, T)],
    ) -> Vec<Vec<(usize, T)>> {
        make_undirected_weighted_by(n, edges, |&(u, v, x)| ([u, v], x))
    }
    pub fn array_make_undirected_weighted<T: Copy>(
        n: usize,
        edges: &[([usize; 2], T)],
    ) -> Vec<Vec<(usize, T)>> {
        make_undirected_weighted_by(n, edges, |&([u, v], x)| ([u, v], x))
    }
    pub fn make_undirected_weighted_by<E, T: Copy>(
        n: usize,
        edges: &[E],
        f: impl Fn(&E) -> ([usize; 2], T),
    ) -> Vec<Vec<(usize, T)>> {
        let mut g = vec![Vec::new(); n];
        for ([u, v], x) in edges.iter().map(f) {
            g[u].push((v, x));
            g[v].push((u, x));
        }
        g
    }
    pub fn tuple_make_directed_weighted<T: Copy>(
        n: usize,
        edges: &[(usize, usize, T)],
    ) -> Vec<Vec<(usize, T)>> {
        make_directed_weighted_by(n, edges, |&(u, v, x)| ([u, v], x))
    }
    pub fn array_make_directed_weighted<T: Copy>(
        n: usize,
        edges: &[([usize; 2], T)],
    ) -> Vec<Vec<(usize, T)>> {
        make_directed_weighted_by(n, edges, |&([u, v], x)| ([u, v], x))
    }
    pub fn make_directed_weighted_by<E, T: Copy>(
        n: usize,
        edges: &[E],
        f: impl Fn(&E) -> ([usize; 2], T),
    ) -> Vec<Vec<(usize, T)>> {
        let mut g = vec![Vec::new(); n];
        edges
            .iter()
            .map(f)
            .for_each(|([u, v], w)| g[u].push((v, w)));
        g
    }
}
// }}}
// template {{{
#[cfg(not(feature = "dbg"))]
#[allow(unused_macros)]
#[macro_export]
macro_rules! lg {
    ($($expr:expr),*) => {};
}
#[allow(dead_code)]
mod ngtio {
    mod i {
        pub use self::{
            multi_token::{Leaf, Parser, ParserTuple, RawTuple, Tuple, VecLen},
            token::{Token, Usize1},
        };
        use std::{
            io::{self, BufRead},
            iter,
        };
        pub fn with_stdin() -> Tokenizer<io::BufReader<io::Stdin>> {
            io::BufReader::new(io::stdin()).tokenizer()
        }
        pub fn with_str(src: &str) -> Tokenizer<&[u8]> {
            src.as_bytes().tokenizer()
        }
        pub struct Tokenizer<S: BufRead> {
            queue: Vec<String>, // FIXME: String のみにすると速そうです。
            scanner: S,
        }
        macro_rules! prim_method {
            ($name:ident: $T:ty) => {
                pub fn $name(&mut self) -> $T {
                    <$T>::leaf().parse(self)
                }
            };
            ($name:ident) => {
                prim_method!($name: $name);
            };
        }
        macro_rules! prim_methods {
            ($name:ident: $T:ty; $($rest:tt)*) => {
                prim_method!($name:$T);
                prim_methods!($($rest)*);
            };
            ($name:ident; $($rest:tt)*) => {
                prim_method!($name);
                prim_methods!($($rest)*);
            };
            () => ()
        }
        impl<S: BufRead> Tokenizer<S> {
            pub fn token(&mut self) -> String {
                self.load();
                self.queue.pop().expect("入力が終了したのですが。")
            }
            pub fn new(scanner: S) -> Self {
                Self {
                    queue: Vec::new(),
                    scanner,
                }
            }
            fn load(&mut self) {
                while self.queue.is_empty() {
                    let mut s = String::new();
                    let length = self.scanner.read_line(&mut s).unwrap(); // 入力が UTF-8 でないときにエラーだそうです。
                    if length == 0 {
                        break;
                    }
                    self.queue = s.split_whitespace().rev().map(str::to_owned).collect();
                }
            }
            pub fn skip_line(&mut self) {
                assert!(
                    self.queue.is_empty(),
                    "行の途中で呼ばないでいただきたいです。現在のトークンキュー: {:?}",
                    &self.queue
                );
                self.load();
            }
            pub fn end(&mut self) {
                self.load();
                assert!(self.queue.is_empty(), "入力はまだあります！");
            }
            pub fn parse<T: Token>(&mut self) -> T::Output {
                T::parse(&self.token())
            }
            pub fn parse_collect<T: Token, B>(&mut self, n: usize) -> B
            where
                B: iter::FromIterator<T::Output>,
            {
                iter::repeat_with(|| self.parse::<T>()).take(n).collect()
            }
            pub fn tuple<T: RawTuple>(&mut self) -> <T::LeafTuple as Parser>::Output {
                T::leaf_tuple().parse(self)
            }
            pub fn vec<T: Token>(&mut self, len: usize) -> Vec<T::Output> {
                T::leaf().vec(len).parse(self)
            }
            pub fn vec_tuple<T: RawTuple>(
                &mut self,
                len: usize,
            ) -> Vec<<T::LeafTuple as Parser>::Output> {
                T::leaf_tuple().vec(len).parse(self)
            }
            pub fn vec2<T: Token>(&mut self, height: usize, width: usize) -> Vec<Vec<T::Output>> {
                T::leaf().vec(width).vec(height).parse(self)
            }
            pub fn vec2_tuple<T>(
                &mut self,
                height: usize,
                width: usize,
            ) -> Vec<Vec<<T::LeafTuple as Parser>::Output>>
            where
                T: RawTuple,
            {
                T::leaf_tuple().vec(width).vec(height).parse(self)
            }
            prim_methods! {
                u8; u16; u32; u64; u128; usize;
                i8; i16; i32; i64; i128; isize;
                f32; f64;
                char; string: String;
            }
        }
        mod token {
            use super::multi_token::Leaf;
            use std::{any, fmt, marker, str};
            pub trait Token: Sized {
                type Output;
                fn parse(s: &str) -> Self::Output;
                fn leaf() -> Leaf<Self> {
                    Leaf(marker::PhantomData)
                }
            }
            impl<T> Token for T
            where
                T: str::FromStr,
                <Self as str::FromStr>::Err: fmt::Debug,
            {
                type Output = Self;
                fn parse(s: &str) -> Self::Output {
                    s.parse().unwrap_or_else(|_| {
                        panic!("Parse error!: ({}: {})", s, any::type_name::<Self>(),)
                    })
                }
            }
            pub struct Usize1 {}
            impl Token for Usize1 {
                type Output = usize;
                fn parse(s: &str) -> Self::Output {
                    usize::parse(s)
                        .checked_sub(1)
                        .expect("Parse error! (Zero substruction error of Usize1)")
                }
            }
        }
        mod multi_token {
            use super::{Token, Tokenizer};
            use std::{io::BufRead, iter, marker};
            pub trait Parser: Sized {
                type Output;
                fn parse<S: BufRead>(&self, server: &mut Tokenizer<S>) -> Self::Output;
                fn vec(self, len: usize) -> VecLen<Self> {
                    VecLen { len, elem: self }
                }
            }
            pub struct Leaf<T>(pub(super) marker::PhantomData<T>);
            impl<T: Token> Parser for Leaf<T> {
                type Output = T::Output;
                fn parse<S: BufRead>(&self, server: &mut Tokenizer<S>) -> T::Output {
                    server.parse::<T>()
                }
            }
            pub struct VecLen<T> {
                pub len: usize,
                pub elem: T,
            }
            impl<T: Parser> Parser for VecLen<T> {
                type Output = Vec<T::Output>;
                fn parse<S: BufRead>(&self, server: &mut Tokenizer<S>) -> Self::Output {
                    iter::repeat_with(|| self.elem.parse(server))
                        .take(self.len)
                        .collect()
                }
            }
            pub trait RawTuple {
                type LeafTuple: Parser;
                fn leaf_tuple() -> Self::LeafTuple;
            }
            pub trait ParserTuple {
                type Tuple: Parser;
                fn tuple(self) -> Self::Tuple;
            }
            pub struct Tuple<T>(pub T);
            macro_rules! impl_tuple {
                ($($t:ident: $T:ident),*) => {
                    impl<$($T),*> Parser for Tuple<($($T,)*)>
                        where
                            $($T: Parser,)*
                            {
                                type Output = ($($T::Output,)*);
#[allow(unused_variables)]
                                fn parse<S: BufRead >(&self, server: &mut Tokenizer<S>) -> Self::Output {
                                    match self {
                                        Tuple(($($t,)*)) => {
                                            ($($t.parse(server),)*)
                                        }
                                    }
                                }
                            }
                    impl<$($T: Token),*> RawTuple for ($($T,)*) {
                        type LeafTuple = Tuple<($(Leaf<$T>,)*)>;
                        fn leaf_tuple() -> Self::LeafTuple {
                            Tuple(($($T::leaf(),)*))
                        }
                    }
                    impl<$($T: Parser),*> ParserTuple for ($($T,)*) {
                        type Tuple = Tuple<($($T,)*)>;
                        fn tuple(self) -> Self::Tuple {
                            Tuple(self)
                        }
                    }
                };
            }
            impl_tuple!();
            impl_tuple!(t1: T1);
            impl_tuple!(t1: T1, t2: T2);
            impl_tuple!(t1: T1, t2: T2, t3: T3);
            impl_tuple!(t1: T1, t2: T2, t3: T3, t4: T4);
            impl_tuple!(t1: T1, t2: T2, t3: T3, t4: T4, t5: T5);
            impl_tuple!(t1: T1, t2: T2, t3: T3, t4: T4, t5: T5, t6: T6);
            impl_tuple!(t1: T1, t2: T2, t3: T3, t4: T4, t5: T5, t6: T6, t7: T7);
            impl_tuple!(
                t1: T1,
                t2: T2,
                t3: T3,
                t4: T4,
                t5: T5,
                t6: T6,
                t7: T7,
                t8: T8
            );
        }
        trait Scanner: BufRead + Sized {
            fn tokenizer(self) -> Tokenizer<Self> {
                Tokenizer::new(self)
            }
        }
        impl<R: BufRead> Scanner for R {}
    }
    pub use self::i::{with_stdin, with_str};
    mod prelude {
        pub use super::i::{Parser, ParserTuple, RawTuple, Token, Usize1};
    }
}
// }}}