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

• 問題: No.898 tri-βutree
• ユーザー: wesrivery
• 提出日時: 2019-10-06 19:01:37
• 言語: Rust (1.83.0 + proconio)
• 結果: WA
• 実行時間: -
• コード長: 5,769 bytes
• コンパイル時間: 23,398 ms
• コンパイル使用メモリ: 400,128 KB
• 実行使用メモリ: 55,808 KB
• 最終ジャッジ日時: 2024-10-10 22:32:48
• 合計ジャッジ時間: 24,608 ms
• ジャッジサーバーID （参考情報）: judge1 / judge3

ファイルパターン	結果
sample	AC * 1
other	AC * 1 WA * 20

コンパイルメッセージ

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

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

ソースコード

use std::cmp::Ordering;

#[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: u32, edges: &Vec<Vec<(usize, u32)>>, mut dist: Vec<Option<u32>>) -> Vec<Option<u32>> {
    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, u32)>>, 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) => {
                    if to - from + 1 == n {
                        return (*v, *index);
                    }
                    let boundary = n / 2;
                    if from < boundary && boundary <= to {
                        let right_to = to - boundary;
                        let (lval, lindex) = left.as_ref().query(from, boundary - 1, boundary);
                        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 {
                        return left.as_ref().query(from, to, boundary);
                    } else {
                        let right_from = from - boundary;
                        let right_to = to - boundary;
                        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
    }
}


#[allow(unused_must_use)]
#[allow(unused_variables)]
fn solve() {
    use segment_tree::SegmentTree;

    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, u32);
        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 (q, ) = input!(usize);
    for _ in 0..q {
        let (x, y, z) = input!(usize, usize, usize);

        let (Log(xy_lca, _), _) = tree.query(first_visit[x], first_visit[y], loglen);
        let (Log(yz_lca, _), _) = tree.query(first_visit[y], first_visit[z], loglen);
        let (Log(zx_lca, _), _) = tree.query(first_visit[x], first_visit[z], loglen);

        let xy_dist = dist_from_root[x].unwrap() + dist_from_root[y].unwrap() - 2 * dist_from_root[xy_lca].unwrap();
        let yz_dist = dist_from_root[y].unwrap() + dist_from_root[z].unwrap() - 2 * dist_from_root[yz_lca].unwrap();
        let zx_dist = dist_from_root[z].unwrap() + dist_from_root[x].unwrap() - 2 * dist_from_root[xy_lca].unwrap();

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

fn main() {
    solve();
}