310. Minimum Height Trees の解き方メモ

「N個のノードからなる根なし木が与えられる。各ノードを根として木の高さを求めたとき、その高さが最小になるようなノードをすべて求めなさい」という問題です。

書き方からはピンと来ないかもしれませんが、これはすなわち「木の中心を求めなさい」ということです。ノードvとwの長さが木の直径と等しいとき、vとwを結ぶパスの中央にあるノードが木の中心になることが知られています。

では木の直径をどのように求めるかですが、これもよく知られたものがあります。

1. 任意のノードからもっとも遠いノードを求める。これをvとする
2. vからもっとも遠いノードを求める。これをwとする
3. vとwの距離が直径となる

あるノードからもっとも遠いノードを求めるには幅優先探索を利用します。ここがO(N)で求められるので、全体の計算量もO(N)になります。

参考:

use std::collections::VecDeque;

impl Solution {
    pub fn find_min_height_trees(n: i32, edges: Vec<Vec<i32>>) -> Vec<i32> {
        let n = n as usize;

        let mut graph = vec![vec![]; n];
        for edge in edges {
            let a = edge[0] as usize;
            let b = edge[1] as usize;
            graph[a].push(b);
            graph[b].push(a);
        }

        let path1 = Self::farest(n, &graph, 0);
        let path2 = Self::farest(n, &graph, path1[0]);

        let m = path2.len();
        if m % 2 == 0 {
            vec![path2[m / 2 - 1] as i32, path2[m / 2] as i32]
        } else {
            vec![path2[m / 2] as i32]
        }
    }

    /// startから最も遠い地点までのパスを求める
    pub fn farest(n: usize, graph: &Vec<Vec<usize>>, start: usize) -> Vec<usize> {
        let mut queue = VecDeque::new();
        let mut distance = vec![-1; n];
        let mut pre = (0..n).into_iter().collect::<Vec<_>>();

        queue.push_back(start);
        distance[start] = 0;

        while let Some(cur) = queue.pop_front() {
            for &nxt in &graph[cur] {
                if distance[nxt] == -1 {
                    distance[nxt] = distance[cur] + 1;
                    pre[nxt] = cur;
                    queue.push_back(nxt);
                }
            }
        }

        let mut path = vec![(0..n).max_by_key(|&i| distance[i]).unwrap()];
        while let Some(&cur) = path.last() {
            if cur == start {
                break;
            }
            path.push(pre[cur]);
        }

        path
    }
}