点击获取AI摘要

3372. 连接两棵树后最大目标节点数目 I M

有两棵无向树，分别有 n 和 m 个树节点。两棵树中的节点编号分别为[0, n - 1] 和 [0, m - 1] 中的整数。

给你两个二维整数 edges1 和 edges2 ，长度分别为 n - 1 和 m - 1 ，其中 edges1[i] = [aᵢ, bᵢ] 表示第一棵树中节点 ai 和 bi 之间有一条边，edges2[i] = [uᵢ, vᵢ] 表示第二棵树中节点 uᵢ 和 vᵢ 之间有一条边。同时给你一个整数 k 。

如果节点 u 和节点 v 之间路径的边数小于等于 k ，那么我们称节点 u 是节点 v 的 目标节点 。注意，一个节点一定是它自己的 目标节点 。

请你返回一个长度为 n 的整数数组 answer ，answer[i] 表示将第一棵树中的一个节点与第二棵树中的一个节点连接一条边后，第一棵树中节点 i 的 目标节点 数目的 最大值 。

注意，每个查询相互独立。意味着进行下一次查询之前，你需要先把刚添加的边给删掉。

示例 1：

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

输出： [9,7,9,8,8]

解释：

对于 i = 0 ，连接第一棵树中的节点 0 和第二棵树中的节点 0 。

对于 i = 1 ，连接第一棵树中的节点 1 和第二棵树中的节点 0 。

对于 i = 2 ，连接第一棵树中的节点 2 和第二棵树中的节点 4 。

对于 i = 3 ，连接第一棵树中的节点 3 和第二棵树中的节点 4 。

对于 i = 4 ，连接第一棵树中的节点 4 和第二棵树中的节点 4 。

示例1

示例 2：

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

输出： [6,3,3,3,3]

解释：

对于每个 i ，连接第一棵树中的节点 i 和第二棵树中的任意一个节点。

示例2

提示：

2 <= n, m <= 1000
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
输入保证 edges1 和 edges2 都表示合法的树。
0 <= k <= 1000

问题分析

给定两棵无向树：
- 第一棵树有 $n$ 个节点，编号 $[0, n-1]$ ，由 edges1 长度为 $n-1$ 的边列表表示；
- 第二棵树有 $m$ 个节点，编号 $[0, m-1]$ ，由 edges2 长度为 $m-1$ 的边列表表示。
定义：对于两棵树中任何两个节点 $u$ 和 $v$ ，如果它们之间的路径边数 $\mathrm{dist}(u,v) \le k$ ，则称 $u$ 是 $v$ 的“目标节点”。
额外操作：对第一棵树中的某个节点 $i$ ，与第二棵树中的某个节点 $j$ 连一条新边（长度记为 1）。然后整个图就是两棵树通过这条新边连接成一个整体。
需要求：对于第一棵树中每个 $i$ ，尝试将 $i$ 与第二棵树中某个 $j$ 连边后，统计“目标节点”的总个数，取最大值。输出长度为 $n$ 的数组 answer，其中 answer[i] 为节点 $i$ 能获得的最大“目标节点”数量。
注意：每次查询相互独立，添加完这条边后要把它删掉再进行下一次查询。

“目标节点”计数可以分为两部分：
- 第一棵树内的“目标节点”：只考虑第一棵树原有的结构，不受连接第二棵树的影响。
- 跨树到第二棵树的“目标节点”：路径必须经过新边，即先从 $i$ 到 $i$ （零步），走新边到 $j$ （1 步），再在第二棵树中走若干步到 $v$ ；总步数 $\le k$ 等价于第二棵树中 $\mathrm{dist}_2(j,v) \le k-1$ 。
因此，对每个 $i$ ：
- 第一部分：统计在第一棵树里，以 $i$ 为起点，所有 $\mathrm{dist}_1(i,u)\le k$ 的节点个数，记为 $C_1[i]$ 。
- 第二部分：选一个在第二棵树的连接点 $j$ ，使得以 $j$ 为起点，在第二棵树内 $\mathrm{dist}_2(j,v)\le k-1$ 的节点数最多。记第二棵树上每个节点到半径 $(k-1)$ 内的可达节点数量为 $C_2[j]$ ，只需取 $\max_j C_2[j]$ 。
- 故 $\displaystyle answer[i] = C_1[i] + \max_{j} C_2[j]$ 。

算法思路

推断节点个数
- 第一棵树节点数 $n = \text{len(edges1)} + 1$ ；
- 第二棵树节点数 $m = \text{len(edges2)} + 1$ 。
构造邻接表
- 用 g1 表示第一棵树的邻接表，大小为 $n$ ；
- 用 g2 表示第二棵树的邻接表，大小为 $m$ 。
BFS 统计“半径内可达”节点数
定义函数 bfs_count(adj, radius)：
- 输入：邻接表 adj（节点数为 $N$ ），半径 radius；
- 对每个节点 $i$ 从 $i$ 出发做一次 BFS，维护一个长度为 $N$ 的距离数组 dist，初始 dist[i]=0 其余 -1；
- 当从队列中弹出节点 $u$ 时，如果 dist[u] > radius 就跳过不扩展；否则把 $u$ 计入计数器，再将所有未访问邻居 $w$ （dist[w]==-1）标记 dist[w]=dist[u]+1 并入队；
- 最终统计以 $i$ 为起点满足距离 $\le$ radius 的节点总数，存入数组 counts[i]；
- 返回长度为 $N$ 的列表 counts。
计算两棵树的 $C_1,C_2$
- 对第一棵树调用 bfs_count(g1, k)，得到 $C_1[i] = |\{u: \mathrm{dist}_1(i,u)\le k\}|$ ；
- 对第二棵树调用 bfs_count(g2, k-1)，得到 $C_2[j] = |\{v: \mathrm{dist}_2(j,v)\le k-1\}|$ $C_{2} [j] = ∣ {v : dist_{2} (j, v) \leq k - 1} ∣$ 。
  - 特殊情况：当 $k=0$ 时，radius = k-1 = -1。在实现里，只有 dist[j]=0 时才满足 dist[j] \le -1 不成立，因此只能把 $j` 自己计入一次，符合“节点一定算自己一个目标节点”的要求。
组合答案
- 先求 max_C2 = max(C2)；
- 然后对每个 $i$ ，answer[i] = C1[i] + max_C2。

时间复杂度

记第一棵树节点数为 $n$ ，第二棵树节点数为 $m$ 。
bfs_count(g1, k) 需要对 $n$ 个节点分别做一次 BFS，每次遍历整棵树 $O(n)$ ，共 $O(n^2)$ 。
bfs_count(g2, k-1) 共 $O(m^2)$ 。
计算最大值及拼答案线性扫描均为 $O(n + m)$ 。
因此总时间复杂度为 $O(n^2 + m^2)$ 。
在题目约束 $n,m \le 1000$ 下， $n^2+m^2 \le 2\times10^6$ ，可以在常规硬件环境下接受。

代码分解

类与方法签名

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

直接接收 edges1, edges2, k，内部推断 $n,m$ 。

推断节点数与构造邻接表

n = len(edges1) + 1
m = len(edges2) + 1

g1 = [[] for _ in range(n)]
for u,v in edges1:
    g1[u].append(v)
    g1[v].append(u)

g2 = [[] for _ in range(m)]
for u,v in edges2:
    g2[u].append(v)
    g2[v].append(u)

BFS 统计函数 bfs_count

def bfs_count(adj: List[List[int]], radius: int) -> List[int]:
    sz = len(adj)
    counts = [0] * sz
    from collections import deque

    for i in range(sz):
        dist = [-1] * sz
        dist[i] = 0
        dq = deque([i])
        cnt = 0
        while dq:
            u = dq.popleft()
            if dist[u] > radius:
                continue
            cnt += 1
            for w in adj[u]:
                if dist[w] == -1:
                    dist[w] = dist[u] + 1
                    dq.append(w)
        counts[i] = cnt
    return counts

对每个 i 初始化距离数组 dist，执行典型的 BFS。
如果 dist[u] > radius 就不再向下扩散。
统计所有满足 dist[u] <= radius 的节点总数。

计算 $C_1, C_2$ 并组合
1
2
3
4
5
6
C1 = bfs_count(g1, k)
C2 = bfs_count(g2, k-1)
max_C2 = max(C2) if C2 else 0

answer = [C1_i + max_C2 for C1_i in C1]
return answer
- 第一棵树用半径 $k$ ，第二棵树用半径 $k-1$ 。
- 取第二个列表的最大值加到 $C_1[i]$ 上。

代码实现

from collections import deque
from typing import List

class Solution:
    def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]], k: int) -> List[int]:
        # 1. 推断两棵树的节点数
        n = len(edges1) + 1
        m = len(edges2) + 1
        
        # 2. 构建邻接表
        g1 = [[] for _ in range(n)]
        for u, v in edges1:
            g1[u].append(v)
            g1[v].append(u)
        
        g2 = [[] for _ in range(m)]
        for u, v in edges2:
            g2[u].append(v)
            g2[v].append(u)
        
        # 3. 定义 BFS 统计半径内可达节点数的函数
        def bfs_count(adj: List[List[int]], radius: int) -> List[int]:
            """
            对邻接表 adj 上每个节点做一次 BFS，统计到距离 <= radius 的节点数目
            返回列表 counts，counts[i] 表示以 i 为起点时距离 <= radius 的节点总数
            """
            sz = len(adj)
            counts = [0] * sz
            for i in range(sz):
                dist = [-1] * sz
                dist[i] = 0
                dq = deque([i])
                cnt = 0
                while dq:
                    u = dq.popleft()
                    # 如果当前距离已经超过 radius，则跳过扩展
                    if dist[u] > radius:
                        continue
                    # 计数当前节点
                    cnt += 1
                    for w in adj[u]:
                        if dist[w] == -1:
                            dist[w] = dist[u] + 1
                            dq.append(w)
                counts[i] = cnt
            return counts
        
        # 4. 分别计算两棵树上每个节点在对应半径内的可达节点数
        C1 = bfs_count(g1, k)      # 第一棵树用半径 k
        C2 = bfs_count(g2, k - 1)  # 第二棵树用半径 k-1（k=0 时恰好只能自己）
        
        # 5. 取第二棵树的最大贡献
        max_C2 = max(C2) if C2 else 0
        
        # 6. 组合答案：answer[i] = C1[i] + max_C2
        answer = [c1 + max_C2 for c1 in C1]
        return answer