結果

問題 No.2634 Tree Distance 3
ユーザー akakimidoriakakimidori
提出日時 2024-02-16 22:12:57
言語 Rust
(1.77.0 + proconio)
結果
AC  
実行時間 402 ms / 3,000 ms
コード長 9,941 bytes
コンパイル時間 11,786 ms
コンパイル使用メモリ 402,116 KB
実行使用メモリ 50,692 KB
最終ジャッジ日時 2024-09-28 20:48:06
合計ジャッジ時間 29,099 ms
ジャッジサーバーID
(参考情報)
judge4 / judge1
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 336 ms
48,992 KB
testcase_01 AC 350 ms
48,968 KB
testcase_02 AC 402 ms
48,944 KB
testcase_03 AC 396 ms
49,088 KB
testcase_04 AC 339 ms
49,012 KB
testcase_05 AC 270 ms
33,080 KB
testcase_06 AC 282 ms
33,072 KB
testcase_07 AC 340 ms
33,020 KB
testcase_08 AC 290 ms
50,508 KB
testcase_09 AC 249 ms
50,480 KB
testcase_10 AC 218 ms
50,648 KB
testcase_11 AC 219 ms
50,692 KB
testcase_12 AC 220 ms
50,488 KB
testcase_13 AC 129 ms
34,576 KB
testcase_14 AC 124 ms
34,772 KB
testcase_15 AC 127 ms
34,772 KB
testcase_16 AC 171 ms
40,860 KB
testcase_17 AC 88 ms
25,052 KB
testcase_18 AC 121 ms
31,104 KB
testcase_19 AC 145 ms
34,824 KB
testcase_20 AC 40 ms
13,300 KB
testcase_21 AC 122 ms
34,548 KB
testcase_22 AC 131 ms
34,292 KB
testcase_23 AC 133 ms
35,548 KB
testcase_24 AC 285 ms
50,124 KB
testcase_25 AC 231 ms
50,116 KB
testcase_26 AC 225 ms
50,028 KB
testcase_27 AC 123 ms
35,124 KB
testcase_28 AC 137 ms
35,092 KB
testcase_29 AC 138 ms
35,176 KB
testcase_30 AC 183 ms
33,128 KB
testcase_31 AC 191 ms
33,240 KB
testcase_32 AC 180 ms
33,320 KB
testcase_33 AC 112 ms
24,184 KB
testcase_34 AC 21 ms
6,820 KB
testcase_35 AC 67 ms
15,444 KB
testcase_36 AC 35 ms
9,820 KB
testcase_37 AC 79 ms
17,656 KB
testcase_38 AC 2 ms
6,816 KB
testcase_39 AC 2 ms
6,820 KB
testcase_40 AC 1 ms
6,820 KB
testcase_41 AC 2 ms
6,820 KB
testcase_42 AC 2 ms
6,820 KB
testcase_43 AC 75 ms
17,464 KB
testcase_44 AC 46 ms
12,372 KB
testcase_45 AC 228 ms
41,836 KB
testcase_46 AC 131 ms
28,776 KB
testcase_47 AC 258 ms
46,696 KB
testcase_48 AC 285 ms
49,252 KB
testcase_49 AC 267 ms
49,160 KB
testcase_50 AC 280 ms
49,252 KB
testcase_51 AC 286 ms
49,300 KB
testcase_52 AC 262 ms
49,276 KB
testcase_53 AC 2 ms
6,816 KB
testcase_54 AC 1 ms
6,820 KB
testcase_55 AC 1 ms
6,820 KB
testcase_56 AC 2 ms
6,820 KB
testcase_57 AC 2 ms
6,820 KB
testcase_58 AC 1 ms
6,816 KB
testcase_59 AC 1 ms
6,816 KB
testcase_60 AC 350 ms
49,560 KB
testcase_61 AC 330 ms
49,432 KB
testcase_62 AC 321 ms
49,580 KB
testcase_63 AC 207 ms
49,120 KB
testcase_64 AC 126 ms
33,540 KB
testcase_65 AC 190 ms
46,260 KB
testcase_66 AC 46 ms
15,820 KB
testcase_67 AC 41 ms
13,396 KB
testcase_68 AC 124 ms
37,604 KB
testcase_69 AC 120 ms
36,532 KB
testcase_70 AC 119 ms
36,632 KB
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: unused import: `std::io::Write`
 --> src/main.rs:2:5
  |
2 | use std::io::Write;
  |     ^^^^^^^^^^^^^^
  |
  = note: `#[warn(unused_imports)]` on by default

warning: type alias `Set` is never used
 --> src/main.rs:5:6
  |
5 | type Set<T> = BTreeSet<T>;
  |      ^^^
  |
  = note: `#[warn(dead_code)]` on by default

warning: type alias `Deque` is never used
 --> src/main.rs:6:6
  |
6 | type Deque<T> = VecDeque<T>;
  |      ^^^^^

ソースコード

diff #

use std::collections::*;
use std::io::Write;

type Map<K, V> = BTreeMap<K, V>;
type Set<T> = BTreeSet<T>;
type Deque<T> = VecDeque<T>;

fn main() {
    input! {
        n: usize,
        a: [u32; n],
        e: [(usize1, usize1); n - 1],
    }
    let mut hld = HLD::new(n);
    for (a, b) in e {
        hld.add_edge(a, b);
    }
    hld.build(0);
    let mut up = vec![];
    let mut down = vec![];
    let mut dist = |a: usize, b: usize| -> usize {
        hld.path(a, b, &mut up, &mut down);
        up.iter()
            .chain(down.iter())
            .map(|p| p.1 - p.0)
            .sum::<usize>()
            - 1
    };
    let mut map = Map::new();
    for (i, a) in a.iter().enumerate() {
        map.entry(*a).or_insert(vec![]).push(i);
    }
    let mut dp = (n, n, 0);
    let mut ans = vec![0; n];
    for (_, x) in map.into_iter().rev() {
        for &x in x.iter() {
            if dp.0 == n {
                dp = (x, x, 0);
            } else if dp.0 == dp.1 {
                dp.1 = x;
                dp.2 = dist(dp.0, dp.1);
            } else {
                let p = dist(dp.0, x);
                let q = dist(dp.1, x);
                if p >= q && p > dp.2 {
                    dp.1 = x;
                    dp.2 = p;
                } else if q >= p && q > dp.2 {
                    dp.0 = x;
                    dp.2 = q;
                }
            }
        }
        for &x in x.iter() {
            ans[x] = dist(dp.0, x).max(dist(dp.1, x));
        }
    }
    use util::*;
    println!("{}", ans.iter().join(" "));
}

// ---------- 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 ----------
// ---------- 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() {
                    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(
        &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] {
            down.push((index[x], index[y] + 1));
        } else {
            up.push((index[y], index[x] + 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]
    }
    pub fn jump(
        &self,
        s: usize,
        t: usize,
        mut k: usize,
        up: &mut Vec<(usize, usize)>,
        down: &mut Vec<(usize, usize)>,
    ) -> Option<usize> {
        assert!(s.max(t) < self.size);
        self.path(s, t, up, down);
        for (l, r) in up.drain(..) {
            if k < r - l {
                return Some(self.vertex(r - 1 - k));
            }
            k -= r - l;
        }
        for (l, r) in down.drain(..) {
            if k < r - l {
                return Some(self.vertex(l + k));
            }
            k -= r - l;
        }
        None
    }
}
// ---------- end Heavy-Light decomposition ----------
mod util {
    pub trait Join {
        fn join(self, sep: &str) -> String;
    }

    impl<T, I> Join for I
    where
        I: Iterator<Item = T>,
        T: std::fmt::Display,
    {
        fn join(self, sep: &str) -> String {
            let mut s = String::new();
            use std::fmt::*;
            for (i, v) in self.enumerate() {
                if i > 0 {
                    write!(&mut s, "{}", sep).ok();
                }
                write!(&mut s, "{}", v).ok();
            }
            s
        }
    }
}
0