#387431 (Rust) No.898 tri-βutree

• 問題: No.898 tri-βutree
• ユーザー: wesrivery
• 提出日時: 2019-10-06 20:11:25
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 550 ms / 4,000 ms
• コード長: 6,201 bytes
• コンパイル時間: 5,115 ms
• コンパイル使用メモリ: 158,232 KB
• 実行使用メモリ: 57,344 KB
• 最終ジャッジ日時: 2024-04-26 08:57:47
• 合計ジャッジ時間: 11,814 ms
• ジャッジサーバーID （参考情報）: judge5 / judge1

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	287 ms; 57,344 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	1 ms; 5,376 KB
testcase_05	AC	1 ms; 5,376 KB
testcase_06	AC	2 ms; 5,376 KB
testcase_07	AC	534 ms; 39,040 KB
testcase_08	AC	534 ms; 39,040 KB
testcase_09	AC	531 ms; 39,040 KB
testcase_10	AC	523 ms; 39,168 KB
testcase_11	AC	505 ms; 38,996 KB
testcase_12	AC	526 ms; 39,168 KB
testcase_13	AC	527 ms; 39,040 KB
testcase_14	AC	519 ms; 39,040 KB
testcase_15	AC	527 ms; 39,040 KB
testcase_16	AC	531 ms; 39,040 KB
testcase_17	AC	515 ms; 39,040 KB
testcase_18	AC	525 ms; 39,040 KB
testcase_19	AC	518 ms; 39,040 KB
testcase_20	AC	529 ms; 39,040 KB
testcase_21	AC	550 ms; 39,040 KB

コンパイルメッセージ

warning: variable does not need to be mutable
   --> main.rs:157:9
    |
157 |     let mut log = vec![];
    |         ----^^^
    |         |
    |         help: remove this `mut`
    |
    = note: `#[warn(unused_mut)]` on by default

warning: variable does not need to be mutable
   --> main.rs:158:9
    |
158 |     let mut visited = vec![None; n];
    |         ----^^^^^^^
    |         |
    |         help: remove this `mut`

warning: 2 warnings emitted

ソースコード

use std::cmp::Ordering;
use segment_tree::SegmentTree;

#[allow(unused_macros)]
macro_rules! input {
    ( $($t:ty),* ) => {{
        let mut s = String::new();
        std::io::stdin().read_line(&mut s);
        let mut splits = s.trim().split_whitespace();
        ($( { splits.next().unwrap().parse::<$t>().unwrap() },)*)
    }}
}

fn dfs(i: usize, weight: u64, edges: &Vec<Vec<(usize, u64)>>, mut dist: Vec<Option<u64>>) -> Vec<Option<u64>> {
    for (to, w) in &edges[i] {
        let total_weight = w + weight;
        if let Some(_) = dist[*to] {
            continue;
        }
        dist[*to] = Some(total_weight);
        dist = dfs(*to, total_weight, edges, dist);
    }
    
    return dist;
}

fn euler_tour(i: usize, depth: u32, edges: &Vec<Vec<(usize, u64)>>, mut log: Vec<Log>, mut visited: Vec<Option<usize>>) -> (Vec<Option<usize>>, Vec<Log>) {
    visited[i] = Some(log.len());
    log.push(Log(i, depth));
    for (to, _) in &edges[i] {
        if let Some(_) = visited[*to] {
            continue;
        }
        let res = euler_tour(*to, depth + 1, edges, log, visited);
        visited = res.0;
        log = res.1;
        log.push(Log(i, depth));
    }
    
    return (visited, log);
}

#[allow(dead_code)]
mod segment_tree {
    pub enum SegmentTree<T: Copy + Ord> {
        Leaf(T),
        Node(T, usize, Box<SegmentTree<T>>, Box<SegmentTree<T>>),
    }
    impl<T: Copy + Ord> SegmentTree<T> {
        pub fn query(&self, from: usize, to: usize, n: usize) -> (T, usize) {
            match self {
                SegmentTree::Leaf(v) => (*v, 0),
                SegmentTree::Node(v, index, left, right) => {
                    // println!("l = {}, r = {}, n = {}", from, to, n);
                    if to - from + 1 == n {
                        return (*v, *index);
                    }
                    let boundary = n / 2;
                    if from < boundary && boundary <= to {
                        let right_to = to - boundary;
                        // println!("  before go left");
                        let (lval, lindex) = left.as_ref().query(from, boundary - 1, boundary);
                        // println!("  before go right");
                        let (rval, rindex) = right.as_ref().query(0, right_to, n - boundary);
                        let rindex = rindex + boundary;
                        if lval < rval {
                            return (lval, lindex);
                        } else {
                            return (rval, rindex);
                        }
                    }
                    if from < boundary {
                        // println!("  before go only left");
                        return left.as_ref().query(from, to, boundary);
                    } else {
                        let right_from = from - boundary;
                        let right_to = to - boundary;
                        // println!("  before go right");
                        let (rval, rindex) =
                            right.as_ref().query(right_from, right_to, n - boundary);
                        return (rval, rindex + boundary);
                    }
                }
            }
        }
        pub fn new(a: &[T], n: usize) -> SegmentTree<T> {
            if n == 1 {
                SegmentTree::Leaf(a[0])
            } else {
                let half = n / 2;
                let left = SegmentTree::new(&a[0..half], half);
                let right = SegmentTree::new(&a[half..n], n - half);
                let (i, val) = if left.val() <= right.val() {
                    (left.ind(), left.val())
                } else {
                    (right.ind() + half, right.val())
                };
                SegmentTree::Node(val, i, Box::new(left), Box::new(right))
            }
        }
        fn val(&self) -> T {
            match *self {
                SegmentTree::Node(v, _, _, _) => v,
                SegmentTree::Leaf(v) => v,
            }
        }
        fn ind(&self) -> usize {
            match *self {
                SegmentTree::Node(_, i, _, _) => i,
                SegmentTree::Leaf(_) => 0,
            }
        }
    }
}

#[derive(Clone, Copy, Eq)]
struct Log(usize, u32);

impl Ord for Log {
    fn cmp(&self, other: &Self) -> Ordering {
        self.1.cmp(&other.1)
    }
}

impl PartialOrd for Log {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for Log {
    fn eq(&self, other: &Self) -> bool {
        self.1 == other.1
    }
}

fn calc_dist(a: usize, b: usize, dist: &Vec<u64>, visited: &Vec<usize>, tree: &SegmentTree<Log>, loglen: usize) -> u64 {
    let x = std::cmp::min(a, b);
    let y = std::cmp::max(a, b);

    let xi = std::cmp::min(visited[x], visited[y]);
    let yi = std::cmp::max(visited[x], visited[y]);

    let (Log(lca, _), _) = tree.query(xi, yi + 1, loglen);
    let xy_dist = dist[x] + dist[y] - 2 * dist[lca];

    return xy_dist;
}

#[allow(unused_must_use)]
#[allow(unused_variables)]
fn solve() {

    let (n, ) = input!(usize);
    let mut edges = vec![vec![]; n];
    let mut dist_from_root = vec![None; n];
    let mut log = vec![];
    let mut visited = vec![None; n];

    for _ in 0..(n-1) {
        let (u, v, w) = input!(usize, usize, u64);
        edges[u].push((v, w));
        edges[v].push((u, w));
    }

    dist_from_root[0] = Some(0);
    dist_from_root = dfs(0, 0, &edges, dist_from_root);

    let (visited, log) = euler_tour(0, 0, &edges, log, visited);

    let loglen = log.len();
    
    let tree = SegmentTree::new(&log[..], loglen);

    let first_visit: Vec<usize> = visited.iter().map(|i| i.unwrap()).collect();
    let dist: Vec<u64> = dist_from_root.iter().map(|i| i.unwrap()).collect();

    let (q, ) = input!(usize);
    for _ in 0..q {
        let (x, y, z) = input!(usize, usize, usize);
        
        let xy_dist = calc_dist(x, y, &dist, &first_visit, &tree, loglen);
        let yz_dist = calc_dist(y, z, &dist, &first_visit, &tree, loglen);
        let zx_dist = calc_dist(z, x, &dist, &first_visit, &tree, loglen);

        println!("{}", (xy_dist + yz_dist + zx_dist) / 2);
    }
}

fn main() {
    solve();
}