3203. 合并两棵树后的最小直径

来源第 404 场周赛 Q4难度困难分数2266

题目描述

给你两棵无向树，分别有 n 和 m 个节点，节点编号分别为 0 到 n - 1 和 0 到 m - 1 。给你两个二维整数数组 edges1 和 edges2 ，长度分别为 n - 1 和 m - 1 ，其中 edges1[i] = [a_i, b_i] 表示在第一棵树中节点 a_i 和 b_i 之间有一条边，edges2[i] = [u_i, v_i] 表示在第二棵树中节点 u_i 和 v_i 之间有一条边。

你必须在第一棵树和第二棵树中分别选一个节点，并用一条边连接它们。

请你返回添加边后得到的树中，最小直径 为多少。

一棵树的直径指的是树中任意两个节点之间的最长路径长度。

示例 1：

输入：edges1 = [[0,1],[0,2],[0,3]], edges2 = [[0,1]]

输出：3

解释：

将第一棵树中的节点 0 与第二棵树中的任意节点连接，得到一棵直径为 3 的树。

示例 2：

输入：edges1 = [[0,1],[0,2],[0,3],[2,4],[2,5],[3,6],[2,7]], edges2 = [[0,1],[0,2],[0,3],[2,4],[2,5],[3,6],[2,7]]

输出：5

解释：

将第一棵树中的节点 0 和第二棵树中的节点 0 连接，可以得到一棵直径为 5 的树。

提示：

1 <= n, m <= 10⁵
edges1.length == n - 1
edges2.length == m - 1
edges1[i].length == edges2[i].length == 2
edges1[i] = [a_i, b_i]
0 <= a_i, b_i < n
edges2[i] = [u_i, v_i]
0 <= u_i, v_i < m
输入保证 edges1 和 edges2 分别表示一棵合法的树。

解法

方法一：两次 DFS

思考

需要在两棵树之间加一条边，使新树直径最小。 \(n,m\le 10^5\)，若枚举连接端点再求直径，则次数为 \(nm\) 量级，无法接受。

新直径只有两类：仍落在某一棵原树内，即 \(\max(d_1,d_2)\)；或穿过新边，此时最优是把两边的中心附近连起来，长度为两棵树半径之和再加 \(1\)。因此只要求出两棵树的直径。树上任取一点走最远到达 \(a\)，再从 \(a\) 走最远到达 \(b\)， \(a\) 到 \(b\) 即为直径，两次 DFS 即可，总时间线性。

我们记 \(d_1\) 和 \(d_2\) 分别为两棵树的直径，那么合并后的树的直径有以下两种情况：

合并后的树的直径为原始的一棵树的直径，即 \(\max(d_1, d_2)\)；
合并后的树的直径经过原始的两棵树。我们分别计算原始的两棵树的半径 \(r_1 = \lceil \frac{d_1}{2} \rceil\) 和 \(r_2 = \lceil \frac{d_2}{2} \rceil\)，那么合并后的树的直径为 \(r_1 + r_2 + 1\)。

我们取这两种情况的最大值即可。

在计算树的直径时，我们可以使用两次 DFS。首先我们任选一点出发，找到距离该点最远的点，记为 \(a\)。然后从点 \(a\) 出发，找到距离点 \(a\) 最远的点，记为 \(b\)。可以证明，点 \(a\) 和点 \(b\) 之间的路径即为树的直径。

时间复杂度 \(O(n + m)\)，空间复杂度 \(O(n + m)\)。其中 \(n\) 和 \(m\) 分别为两棵树的节点数。

Python3JavaC++GoTypeScriptRustC#

class Solution:
    def minimumDiameterAfterMerge(
        self, edges1: List[List[int]], edges2: List[List[int]]
    ) -> int:
        d1 = self.treeDiameter(edges1)
        d2 = self.treeDiameter(edges2)
        return max(d1, d2, (d1 + 1) // 2 + (d2 + 1) // 2 + 1)

    def treeDiameter(self, edges: List[List[int]]) -> int:
        def dfs(i: int, fa: int, t: int):
            for j in g[i]:
                if j != fa:
                    dfs(j, i, t + 1)
            nonlocal ans, a
            if ans < t:
                ans = t
                a = i

        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = a = 0
        dfs(0, -1, 0)
        dfs(a, -1, 0)
        return ans

class Solution {
    private List<Integer>[] g;
    private int ans;
    private int a;

    public int minimumDiameterAfterMerge(int[][] edges1, int[][] edges2) {
        int d1 = treeDiameter(edges1);
        int d2 = treeDiameter(edges2);
        return Math.max(Math.max(d1, d2), (d1 + 1) / 2 + (d2 + 1) / 2 + 1);
    }

    public int treeDiameter(int[][] edges) {
        int n = edges.length + 1;
        g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        ans = 0;
        a = 0;
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        dfs(0, -1, 0);
        dfs(a, -1, 0);
        return ans;
    }

    private void dfs(int i, int fa, int t) {
        for (int j : g[i]) {
            if (j != fa) {
                dfs(j, i, t + 1);
            }
        }
        if (ans < t) {
            ans = t;
            a = i;
        }
    }
}

class Solution {
public:
    int minimumDiameterAfterMerge(vector<vector<int>>& edges1, vector<vector<int>>& edges2) {
        int d1 = treeDiameter(edges1);
        int d2 = treeDiameter(edges2);
        return max({d1, d2, (d1 + 1) / 2 + (d2 + 1) / 2 + 1});
    }

    int treeDiameter(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        vector<int> g[n];
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        int ans = 0, a = 0;
        auto dfs = [&](this auto&& dfs, int i, int fa, int t) -> void {
            for (int j : g[i]) {
                if (j != fa) {
                    dfs(j, i, t + 1);
                }
            }
            if (ans < t) {
                ans = t;
                a = i;
            }
        };
        dfs(0, -1, 0);
        dfs(a, -1, 0);
        return ans;
    }
};

func minimumDiameterAfterMerge(edges1 [][]int, edges2 [][]int) int {
    d1 := treeDiameter(edges1)
    d2 := treeDiameter(edges2)
    return max(d1, d2, (d1+1)/2+(d2+1)/2+1)
}

func treeDiameter(edges [][]int) (ans int) {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    a := 0
    var dfs func(i, fa, t int)
    dfs = func(i, fa, t int) {
        for _, j := range g[i] {
            if j != fa {
                dfs(j, i, t+1)
            }
        }
        if ans < t {
            ans = t
            a = i
        }
    }
    dfs(0, -1, 0)
    dfs(a, -1, 0)
    return
}

function minimumDiameterAfterMerge(edges1: number[][], edges2: number[][]): number {
    const d1 = treeDiameter(edges1);
    const d2 = treeDiameter(edges2);
    return Math.max(d1, d2, Math.ceil(d1 / 2) + Math.ceil(d2 / 2) + 1);
}

function treeDiameter(edges: number[][]): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    let [ans, a] = [0, 0];
    const dfs = (i: number, fa: number, t: number): void => {
        for (const j of g[i]) {
            if (j !== fa) {
                dfs(j, i, t + 1);
            }
        }
        if (ans < t) {
            ans = t;
            a = i;
        }
    };
    dfs(0, -1, 0);
    dfs(a, -1, 0);
    return ans;
}

impl Solution {
    pub fn minimum_diameter_after_merge(edges1: Vec<Vec<i32>>, edges2: Vec<Vec<i32>>) -> i32 {
        let d1 = Self::tree_diameter(&edges1);
        let d2 = Self::tree_diameter(&edges2);
        d1.max(d2).max((d1 + 1) / 2 + (d2 + 1) / 2 + 1)
    }

    fn tree_diameter(edges: &Vec<Vec<i32>>) -> i32 {
        let n = edges.len() + 1;
        let mut g = vec![vec![]; n];
        for e in edges {
            let a = e[0] as usize;
            let b = e[1] as usize;
            g[a].push(b);
            g[b].push(a);
        }
        let mut ans = 0;
        let mut a = 0;
        fn dfs(g: &Vec<Vec<usize>>, i: usize, fa: isize, t: i32, ans: &mut i32, a: &mut usize) {
            for &j in &g[i] {
                if j as isize != fa {
                    dfs(g, j, i as isize, t + 1, ans, a);
                }
            }
            if *ans < t {
                *ans = t;
                *a = i;
            }
        }
        dfs(&g, 0, -1, 0, &mut ans, &mut a);
        dfs(&g, a, -1, 0, &mut ans, &mut a);
        ans
    }
}

public class Solution {
    private List<int>[] g;
    private int ans;
    private int a;

    public int MinimumDiameterAfterMerge(int[][] edges1, int[][] edges2) {
        int d1 = TreeDiameter(edges1);
        int d2 = TreeDiameter(edges2);
        return Math.Max(Math.Max(d1, d2), (d1 + 1) / 2 + (d2 + 1) / 2 + 1);
    }

    public int TreeDiameter(int[][] edges) {
        int n = edges.Length + 1;
        g = new List<int>[n];
        for (int k = 0; k < n; ++k) {
            g[k] = new List<int>();
        }
        ans = 0;
        a = 0;
        foreach (var e in edges) {
            int a = e[0], b = e[1];
            g[a].Add(b);
            g[b].Add(a);
        }
        Dfs(0, -1, 0);
        Dfs(a, -1, 0);
        return ans;
    }

    private void Dfs(int i, int fa, int t) {
        foreach (int j in g[i]) {
            if (j != fa) {
                Dfs(j, i, t + 1);
            }
        }
        if (ans < t) {
            ans = t;
            a = i;
        }
    }
}

方法二：两次显式栈

思考

需要在两棵树之间加一条边，使新树直径最小。 \(n,m\le 10^5\)，若枚举连接端点再求直径，则次数为 \(nm\) 量级，无法接受。从任意点递归走到最远点，链上的调用深度就是 \(n\)，会超出递归栈。

新直径只有两类：仍落在某一棵原树内，即 \(\max(d_1,d_2)\)；或穿过新边，此时最优是把两边的中心附近连起来，长度为两棵树半径之和再加 \(1\)。因此只要求出两棵树的直径。

显式栈任取一点走到最远的 \(a\)，再从 \(a\) 走到最远的 \(b\)， \(a\) 到 \(b\) 即为直径。两棵树各走两次，总时间线性。