Leetcode 3313. 查找树中最后标记的节点

1.题目基本信息

1.1.题目描述

有一棵有 n 个节点，编号从 0 到 n - 1 的 无向 树。给定一个长度为 n - 1 的整数数组 edges，其中 edges[i] = [ui, vi] 表示树中的 ui 和 vi 之间有一条边。

一开始，所有 节点都 未标记。之后的每一秒，你需要标记所有 至少 有一个已标记节点相邻的未标记节点。

返回一个数组 nodes，表示在时刻 t = 0 标记了节点 i，那么树中最后标记的节点是 nodes[i]。如果对于任意节点 i 有多个 nodes[i]，你可以选择 任意 一个作为答案。

1.2.题目地址

https://leetcode.cn/problems/find-the-last-marked-nodes-in-tree/description/

2.解题方法

2.1.解题思路

深度优先搜索+广度优先搜索

时间复杂度：O(n)

2.2.解题步骤

第一步，构建无向图的邻接表

第二步，深度优先搜索求树的直径的两个端点p1,p2。参考：Leetcode 1522. N 叉树的直径

2.1.递归出口；当node为叶节点时，递归退出（注意排除根节点0）

2.2.递归主体；使用maxLen1和maxLen2分别维护node结点到达叶节点第一长和第二长的路径长度，node1和node2维护node结点到达叶节点第一长和第二长的路径的叶节点

2.3.调整顺序，保证subMaxLen1>=subMaxLen2

2.4.根据子节点的到叶节点的最长路径更新当前节点到达叶节点的第一长和第二长路径信息

2.5.更新直径和直径两端节点的信息

第三步，调用深搜递归函数，更新直径信息。随后分别对p1,p2两个结点进行广度优先搜索，获取各个结点距离端点p1,p2的距离的数组dists1,dists2

第四步，遍历所有节点，枚举结点node，根据dists1和dists2选择p1和p2中距离node更远的节点，构建数组result，并进行返回

3.解题代码

Python代码

from collections import deque

class Solution:
    def lastMarkedNodes(self, edges: List[List[int]]) -> List[int]:
        # 思路：深度优先搜索+广度优先搜索
        n = len(edges) + 1  # 结点的总数
        # 第一步，构建无向图的邻接表
        graph = [[] for _ in range(n)]
        for a, b in edges:
            graph[a].append(b)
            graph[b].append(a)
        # 第二步，深度优先搜索求树的直径的两个端点p1,p2。参考：Leetcode 1522. N 叉树的直径
        visited = set()
        self.maxLength, self.p1, self.p2 = 0, 0, 0
        def dfs(node:int) -> list[int]:
            visited.add(node)
            # 2.1.递归出口；当node为叶节点时，递归退出（注意排除根节点0）
            if node != 0 and len(graph[node]) == 1:
                return 0, 0, node, node
            # 2.2.递归主体；使用maxLen1和maxLen2分别维护node结点到达叶节点第一长和第二长的路径长度，node1和node2维护node结点到达叶节点第一长和第二长的路径的叶节点
            maxLen1, maxLen2, node1, node2 = 0, 0, 0, 0
            for child in graph[node]:
                if child in visited:
                    # 如果已经访问过了，则跳过，模拟N叉树
                    continue
                subMaxLen1, subMaxLen2, subNode1, subNode2 = dfs(child)
                # 2.3.调整顺序，保证subMaxLen1>=subMaxLen2
                if subMaxLen1 < subMaxLen2:
                    subMaxLen1, subMaxLen2 = subMaxLen2, subMaxLen1
                    subNode1, subNode2 = subNode2, subNode1
                # 2.4.根据子节点的到叶节点的最长路径更新当前节点到达叶节点的第一长和第二长路径信息
                if subMaxLen1 + 1 >= maxLen1:
                    maxLen2, node2 = maxLen1, node1
                    maxLen1, node1 = subMaxLen1 + 1, subNode1
                elif subMaxLen1 + 1 >= maxLen2:
                    maxLen2, node2 = subMaxLen1 + 1, subNode1
                # 2.5.更新直径和直径两端节点的信息
                if self.maxLength <= maxLen1 + maxLen2:
                    self.maxLength = maxLen1 + maxLen2
                    self.p1, self.p2 = node1, node2
            return maxLen1, maxLen2, node1, node2
        # print(self.p1, self.p2)
        # 第三步，调用深搜递归函数，更新直径信息。随后分别对p1,p2两个结点进行广度优先搜索，获取各个结点距离端点p1,p2的距离的数组dists1,dists2
        dfs(0)
        dists1, dists2 = [0] * n, [0] * n
        visited1, visited2 = set([self.p1]), set([self.p2])
        que1, que2 = deque([self.p1]), deque([self.p2])
        slice1, slice2 = 1, 1
        while que1:
            for _ in range(len(que1)):
                node = que1.popleft()
                for child in graph[node]:
                    if child not in visited1:
                        que1.append(child)
                        visited1.add(child)
                        dists1[child] = slice1
            slice1 += 1
        # print(dists1)
        while que2:
            for _ in range(len(que2)):
                node = que2.popleft()
                for child in graph[node]:
                    if child not in visited2:
                        que2.append(child)
                        visited2.add(child)
                        dists2[child] = slice2
            slice2 += 1
        # print(dists2)
        # 第四步，遍历所有节点，枚举结点node，根据dists1和dists2选择p1和p2中距离node更远的节点，构建数组result，并进行返回
        result = [0] * n
        for i in range(n):
            result[i] = self.p1 if dists1[i] > dists2[i] else self.p2
        return result

4.执行结果