跳转至

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

题目描述

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

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

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

 

示例 1：

输入：edges = [[0,1],[0,2]]

输出：[2,2,1]

解释：

  • 对于 i = 0，节点以如下序列标记：[0] -> [0,1,2]。1 和 2 都可以是答案。
  • 对于 i = 1，节点以如下序列标记：[1] -> [0,1] -> [0,1,2]。节点 2 最后被标记。
  • 对于 i = 2，节点以如下序列标记：[2] -> [0,2] -> [0,1,2]。节点 1 最后被标记。

示例 2：

输入：edges = [[0,1]]

输出：[1,0]

解释：

  • 对于 i = 0，节点以如下序列被标记：[0] -> [0,1]。
  • 对于 i = 1，节点以如下序列被标记：[1] -> [0,1]。

示例 3：

输入：edges = [[0,1],[0,2],[2,3],[2,4]]

输出：[3,3,1,1,1]

解释：

  • 对于 i = 0，节点以如下序列被标记：[0] -> [0,1,2] -> [0,1,2,3,4]。
  • 对于 i = 1，节点以如下序列被标记：[1] -> [0,1] -> [0,1,2] -> [0,1,2,3,4]。
  • 对于 i = 2，节点以如下序列被标记：[2] -> [0,2,3,4] -> [0,1,2,3,4]。
  • 对于 i = 3，节点以如下序列被标记：[3] -> [2,3] -> [0,2,3,4] -> [0,1,2,3,4]。
  • 对于 i = 4，节点以如下序列被标记：[4] -> [2,4] -> [0,2,3,4] -> [0,1,2,3,4]。

 

提示：

  • 2 <= n <= 105
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= edges[i][0], edges[i][1] <= n - 1
  • 输入保证 edges 形成一棵合法的树。

解法

方法一：求树的直径 + DFS

根据题目描述，最后一个被标记的节点一定是树的直径的一个端点，因为树的直径上的节点到直径上的任意一个节点的距离最大。

我们可以从任意一个节点开始进行深度优先搜索，找到距离最远的节点 $a$，这个节点就是树的直径的一个端点。

然后从节点 $a$ 开始进行深度优先搜索，找到距离最远的节点 $b$，这个节点就是树的直径的另一个端点，在这个过程中，我们计算出了每个节点到节点 $a$ 的距离，记为 $\textit{dist2}$。

接着从节点 $b$ 开始进行深度优先搜索，计算出每个节点到节点 $b$ 的距离，记为 $\textit{dist3}$。

那么，对于每一个节点 $i$，如果 $\textit{dist2}[i] > \textit{dist3}[i]$，那么节点 $a$ 到节点 $i$ 的距离更远，所以节点 $a$ 是最后一个被标记的节点；否则，节点 $b$ 是最后一个被标记的节点。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 是节点的数量。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Solution:
    def lastMarkedNodes(self, edges: List[List[int]]) -> List[int]:
        def dfs(i: int, fa: int, dist: List[int]):
            for j in g[i]:
                if j != fa:
                    dist[j] = dist[i] + 1
                    dfs(j, i, dist)

        n = len(edges) + 1
        g = [[] for _ in range(n)]
        for u, v in edges:
            g[u].append(v)
            g[v].append(u)

        dist1 = [-1] * n
        dist1[0] = 0
        dfs(0, -1, dist1)
        a = dist1.index(max(dist1))

        dist2 = [-1] * n
        dist2[a] = 0
        dfs(a, -1, dist2)
        b = dist2.index(max(dist2))

        dist3 = [-1] * n
        dist3[b] = 0
        dfs(b, -1, dist3)

        return [a if x > y else b for x, y in zip(dist2, dist3)]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
class Solution {
    private List<Integer>[] g;

    public int[] lastMarkedNodes(int[][] edges) {
        int n = edges.length + 1;
        g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (var e : edges) {
            int u = e[0], v = e[1];
            g[u].add(v);
            g[v].add(u);
        }
        int[] dist1 = new int[n];
        dist1[0] = 0;
        dfs(0, -1, dist1);
        int a = maxNode(dist1);

        int[] dist2 = new int[n];
        dist2[a] = 0;
        dfs(a, -1, dist2);
        int b = maxNode(dist2);

        int[] dist3 = new int[n];
        dist3[b] = 0;
        dfs(b, -1, dist3);

        int[] ans = new int[n];
        for (int i = 0; i < n; ++i) {
            ans[i] = dist2[i] > dist3[i] ? a : b;
        }
        return ans;
    }

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

    private int maxNode(int[] dist) {
        int mx = 0;
        for (int i = 0; i < dist.length; ++i) {
            if (dist[mx] < dist[i]) {
                mx = i;
            }
        }
        return mx;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class Solution {
public:
    vector<int> lastMarkedNodes(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        g.resize(n);
        for (const auto& e : edges) {
            int u = e[0], v = e[1];
            g[u].push_back(v);
            g[v].push_back(u);
        }
        vector<int> dist1(n);
        dfs(0, -1, dist1);
        int a = max_element(dist1.begin(), dist1.end()) - dist1.begin();

        vector<int> dist2(n);
        dfs(a, -1, dist2);
        int b = max_element(dist2.begin(), dist2.end()) - dist2.begin();

        vector<int> dist3(n);
        dfs(b, -1, dist3);

        vector<int> ans;
        for (int i = 0; i < n; ++i) {
            ans.push_back(dist2[i] > dist3[i] ? a : b);
        }
        return ans;
    }

private:
    vector<vector<int>> g;

    void dfs(int i, int fa, vector<int>& dist) {
        for (int j : g[i]) {
            if (j != fa) {
                dist[j] = dist[i] + 1;
                dfs(j, i, dist);
            }
        }
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
func lastMarkedNodes(edges [][]int) (ans []int) {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        u, v := e[0], e[1]
        g[u] = append(g[u], v)
        g[v] = append(g[v], u)
    }
    var dfs func(int, int, []int)
    dfs = func(i, fa int, dist []int) {
        for _, j := range g[i] {
            if j != fa {
                dist[j] = dist[i] + 1
                dfs(j, i, dist)
            }
        }
    }
    maxNode := func(dist []int) int {
        mx := 0
        for i, d := range dist {
            if dist[mx] < d {
                mx = i
            }
        }
        return mx
    }

    dist1 := make([]int, n)
    dfs(0, -1, dist1)
    a := maxNode(dist1)

    dist2 := make([]int, n)
    dfs(a, -1, dist2)
    b := maxNode(dist2)

    dist3 := make([]int, n)
    dfs(b, -1, dist3)

    for i, x := range dist2 {
        if x > dist3[i] {
            ans = append(ans, a)
        } else {
            ans = append(ans, b)
        }
    }
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
function lastMarkedNodes(edges: number[][]): number[] {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [u, v] of edges) {
        g[u].push(v);
        g[v].push(u);
    }
    const dfs = (i: number, fa: number, dist: number[]) => {
        for (const j of g[i]) {
            if (j !== fa) {
                dist[j] = dist[i] + 1;
                dfs(j, i, dist);
            }
        }
    };

    const dist1: number[] = Array(n).fill(0);
    dfs(0, -1, dist1);
    const a = dist1.indexOf(Math.max(...dist1));

    const dist2: number[] = Array(n).fill(0);
    dfs(a, -1, dist2);
    const b = dist2.indexOf(Math.max(...dist2));

    const dist3: number[] = Array(n).fill(0);
    dfs(b, -1, dist3);

    const ans: number[] = [];
    for (let i = 0; i < n; ++i) {
        ans.push(dist2[i] > dist3[i] ? a : b);
    }
    return ans;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
/**
 * @param {number[][]} edges
 * @return {number[]}
 */
var lastMarkedNodes = function (edges) {
    const n = edges.length + 1;
    const g = Array.from({ length: n }, () => []);
    for (const [u, v] of edges) {
        g[u].push(v);
        g[v].push(u);
    }
    const dfs = (i, fa, dist) => {
        for (const j of g[i]) {
            if (j !== fa) {
                dist[j] = dist[i] + 1;
                dfs(j, i, dist);
            }
        }
    };

    const dist1 = Array(n).fill(0);
    dfs(0, -1, dist1);
    const a = dist1.indexOf(Math.max(...dist1));

    const dist2 = Array(n).fill(0);
    dfs(a, -1, dist2);
    const b = dist2.indexOf(Math.max(...dist2));

    const dist3 = Array(n).fill(0);
    dfs(b, -1, dist3);

    const ans = [];
    for (let i = 0; i < n; ++i) {
        ans.push(dist2[i] > dist3[i] ? a : b);
    }
    return ans;
};

评论