3203. Find Minimum Diameter After Merging Two Trees

Problem

There exist two undirected trees with n and m nodes, numbered from 0 to n - 1 and from 0 to m - 1, respectively. You are given two 2D integer arrays edges1 and edges2 of lengths n - 1 and m - 1, respectively, where edges1[i] = [aᵢ, bᵢ] indicates that there is an edge between nodes aᵢ and bᵢ in the first tree and edges2[i] = [uᵢ, vᵢ] indicates that there is an edge between nodes uᵢ and vᵢ in the second tree.

You must connect one node from the first tree with another node from the second tree with an edge.

Return the minimum possible diameter of the resulting tree.

The diameter of a tree is the length of the longest path between any two nodes in the tree.

https://leetcode.com/problems/find-minimum-diameter-after-merging-two-trees/

Example 1:

case1

Input: edges1 = [[0,1],[0,2],[0,3]], edges2 = [[0,1]]
Output: 3
Explanation:
We can obtain a tree of diameter 3 by connecting node 0 from the first tree with any node from the second tree.

case2

Example 2:

Input: 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]]
Output: 5
Explanation:
We can obtain a tree of diameter 5 by connecting node 0 from the first tree with node 0 from the second tree.

Constraints:

1 <= n, m <= 10⁵
edges1.length == n - 1
edges2.length == m - 1
edges1[i].length == edges2[i].length == 2
edges1[i] = [aᵢ, bᵢ]
0 <= aᵢ, bᵢ < n
edges2[i] = [uᵢ, vᵢ]
0 <= uᵢ, vᵢ < m
The input is generated such that edges1 and edges2 represent valid trees.

Test Cases

1 2	class Solution: def minimumDiameterAfterMerge(self, edges1: List[List[int]], edges2: List[List[int]]) -> int:

solution_test.py

import pytest

from solution import Solution


@pytest.mark.parametrize('edges1, edges2, expected', [
    ([[0,1],[0,2],[0,3]], [[0,1]], 3),
    ([[0,1],[0,2],[0,3],[2,4],[2,5],[3,6],[2,7]], [[0,1],[0,2],[0,3],[2,4],[2,5],[3,6],[2,7]], 5),

    ([], [], 1),
    ([[0,1],[2,0],[3,2],[3,6],[8,7],[4,8],[5,4],[3,5],[3,9]], [[0,1],[0,2],[0,3]], 7),
])
@pytest.mark.parametrize('sol', [Solution()])
def test_solution(sol, edges1, edges2, expected):
    assert sol.minimumDiameterAfterMerge(edges1, edges2) == expected

Thoughts

在两棵树中分别找到最长的的路径（也就是直径），取正中的点（如果正中有两个点，就任取一个），用新的边连起来，这样构造出来的新树的直径是最小的。

设两棵树的直径分别是 d1 和 d2，那么新树的直径是 min{d1, d2, 1 + ⌈d1/2⌉ + ⌈d2/2⌉}。

求一棵无向树的直径，可以从任一顶点出发，遍历树并找到距离最远的点，再从这个点出发找到距离最远的点以及距离，这个距离就是树的直径。时间复杂度 O(n)。

总的时间复杂度 O(m + n)，空间复杂度 O(m + n)。

Code

solution.py

from collections import defaultdict
from math import ceil


def dfs(tree: dict[int, list[int]], u: int) -> tuple[int, int]:
    """Finds the farthest node from u, and the distance between them."""
    visited = [False] * len(tree)
    farthest = u, 0
    stack = [farthest] # (node, depth)
    while stack:
        u, d = stack.pop()
        if visited[u]: continue
        visited[u] = True
        if d > farthest[1]:
            farthest = u, d

        d += 1
        for v in tree[u]:
            if not visited[v]: stack.append((v, d))

    return farthest


def diameter(edges: list[list[int]]) -> int:
    if not edges: return 0
    tree: dict[int, list[int]] = defaultdict(list)
    for u, v in edges:
        tree[u].append(v)
        tree[v].append(u)

    u, _ = dfs(tree, 0)
    _, d = dfs(tree, u)
    return d


class Solution:
    def minimumDiameterAfterMerge(self, edges1: list[list[int]], edges2: list[list[int]]) -> int:
        ds = list(map(diameter, (edges1, edges2)))
        return max(ds[0], ds[1], 1 + ceil(ds[0] / 2) + ceil(ds[1] / 2))