跳转至

2049. 统计最高分的节点数目

来源第 264 场周赛 Q3难度中等分数1911

题目描述

给你一棵根节点为 0 的 二叉树 ，它总共有 n 个节点，节点编号为 0 到 n - 1 。同时给你一个下标从 0 开始的整数数组 parents 表示这棵树，其中 parents[i] 是节点 i 的父节点。由于节点 0 是根，所以 parents[0] == -1 。

一个子树的 大小 为这个子树内节点的数目。每个节点都有一个与之关联的 分数 。求出某个节点分数的方法是，将这个节点和与它相连的边全部 删除 ，剩余部分是若干个 非空 子树，这个节点的 分数 为所有这些子树 大小的乘积 。

请你返回有 最高得分 节点的 数目 。

 

示例 1:

example-1

输入：parents = [-1,2,0,2,0]
输出：3
解释：
- 节点 0 的分数为：3 * 1 = 3
- 节点 1 的分数为：4 = 4
- 节点 2 的分数为：1 * 1 * 2 = 2
- 节点 3 的分数为：4 = 4
- 节点 4 的分数为：4 = 4
最高得分为 4 ，有三个节点得分为 4 （分别是节点 1，3 和 4 ）。

示例 2：

example-2

输入：parents = [-1,2,0]
输出：2
解释：
- 节点 0 的分数为：2 = 2
- 节点 1 的分数为：2 = 2
- 节点 2 的分数为：1 * 1 = 1
最高分数为 2 ，有两个节点分数为 2 （分别为节点 0 和 1 ）。

 

提示：

  • n == parents.length
  • 2 <= n <= 105
  • parents[0] == -1
  • 对于 i != 0 ，有 0 <= parents[i] <= n - 1
  • parents 表示一棵二叉树。

解法

方法一：DFS

思考

去掉结点后，分数为各连通块大小之积。 \(n \le 10^5\)，需一次遍历得到所有分数。子树大小在 DFS 回溯时可知，补集大小为 \(n-\textit{cnt}\)。

建出孩子邻接表，从根 DFS：把各孩子子树大小相乘，若还有上方块再乘 \(n-\textit{cnt}\)。同步维护最高分及其出现次数。

我们先根据给定的父节点数组 parents 构建图 \(g\)，其中 \(g[i]\) 表示节点 \(i\) 的所有子节点。定义变量 \(\textit{ans}\) 表示最高得分的节点数目，变量 \(mx\) 表示最高得分。

然后，我们设计一个函数 \(\textit{dfs}(i, fa)\)，它的作用是计算节点 \(i\) 的分数，并返回以节点 \(i\) 为根的子树的节点数目。

函数 \(\textit{dfs}(i, fa)\) 的计算过程如下：

我们首先初始化变量 \(\textit{cnt} = 1\)，表示以节点 \(i\) 为根的子树的节点数目；变量 \(\textit{score} = 1\)，表示以节点 \(i\) 初始分数。

接下来，我们遍历节点 \(i\) 的所有子节点 \(j\)，如果 \(j\) 不是节点 \(i\) 的父节点 \(fa\)，那么我们递归调用 \(\textit{dfs}(j, i)\)，并将返回值累乘到 \(\textit{score}\) 中，同时将返回值累加到 \(\textit{cnt}\) 中。

遍历完子节点后，如果 \(n - \textit{cnt} > 0\)，那么我们将 \(n - \textit{cnt}\) 累乘到 \(\textit{score}\) 中。

然后，我们判断 \(mx\) 是否小于 \(\textit{score}\)，如果小于，那么我们将 \(mx\) 更新为 \(\textit{score}\)，并将 \(\textit{ans}\) 更新为 \(1\)；如果等于，那么我们将 \(\textit{ans}\) 更新为 \(\textit{ans} + 1\)。

最后，我们返回 \(\textit{cnt}\)。

最终，我们调用 \(\textit{dfs}(0, -1)\)，并返回 \(\textit{ans}\)。

时间复杂度 \(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
class Solution:
    def countHighestScoreNodes(self, parents: List[int]) -> int:
        def dfs(i: int, fa: int):
            cnt = score = 1
            for j in g[i]:
                if j != fa:
                    t = dfs(j, i)
                    score *= t
                    cnt += t
            if n - cnt:
                score *= n - cnt
            nonlocal ans, mx
            if mx < score:
                mx = score
                ans = 1
            elif mx == score:
                ans += 1
            return cnt

        n = len(parents)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[parents[i]].append(i)
        ans = mx = 0
        dfs(0, -1)
        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
38
39
class Solution {
    private List<Integer>[] g;
    private int ans;
    private long mx;
    private int n;

    public int countHighestScoreNodes(int[] parents) {
        n = parents.length;
        g = new List[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int i = 1; i < n; ++i) {
            g[parents[i]].add(i);
        }
        dfs(0, -1);
        return ans;
    }

    private int dfs(int i, int fa) {
        int cnt = 1;
        long score = 1;
        for (int j : g[i]) {
            if (j != fa) {
                int t = dfs(j, i);
                cnt += t;
                score *= t;
            }
        }
        if (n - cnt > 0) {
            score *= n - cnt;
        }
        if (mx < score) {
            mx = score;
            ans = 1;
        } else if (mx == score) {
            ++ans;
        }
        return cnt;
    }
}

 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
class Solution {
public:
    int countHighestScoreNodes(vector<int>& parents) {
        int n = parents.size();
        vector<int> g[n];
        for (int i = 1; i < n; ++i) {
            g[parents[i]].push_back(i);
        }
        int ans = 0;
        long long mx = 0;
        function<int(int, int)> dfs = [&](int i, int fa) {
            long long score = 1;
            int cnt = 1;
            for (int j : g[i]) {
                if (j != fa) {
                    int t = dfs(j, i);
                    cnt += t;
                    score *= t;
                }
            }
            if (n - cnt) {
                score *= n - cnt;
            }
            if (mx < score) {
                mx = score;
                ans = 1;
            } else if (mx == score) {
                ++ans;
            }
            return cnt;
        };
        dfs(0, -1);
        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
func countHighestScoreNodes(parents []int) (ans int) {
    n := len(parents)
    g := make([][]int, n)
    for i := 1; i < n; i++ {
        g[parents[i]] = append(g[parents[i]], i)
    }
    mx := 0
    var dfs func(i, fa int) int
    dfs = func(i, fa int) int {
        cnt, score := 1, 1
        for _, j := range g[i] {
            if j != fa {
                t := dfs(j, i)
                cnt += t
                score *= t
            }
        }
        if n-cnt > 0 {
            score *= n - cnt
        }
        if mx < score {
            mx = score
            ans = 1
        } else if mx == score {
            ans++
        }
        return cnt
    }
    dfs(0, -1)
    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
function countHighestScoreNodes(parents: number[]): number {
    const n = parents.length;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (let i = 1; i < n; i++) {
        g[parents[i]].push(i);
    }
    let [ans, mx] = [0, 0];
    const dfs = (i: number, fa: number): number => {
        let [cnt, score] = [1, 1];
        for (const j of g[i]) {
            if (j !== fa) {
                const t = dfs(j, i);
                cnt += t;
                score *= t;
            }
        }
        if (n - cnt) {
            score *= n - cnt;
        }
        if (mx < score) {
            mx = score;
            ans = 1;
        } else if (mx === score) {
            ans++;
        }
        return cnt;
    };
    dfs(0, -1);
    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
38
39
40
41
42
43
44
45
46
public class Solution {
    private List<int>[] g;
    private int ans;
    private long mx;
    private int n;

    public int CountHighestScoreNodes(int[] parents) {
        n = parents.Length;
        g = new List<int>[n];
        for (int i = 0; i < n; ++i) {
            g[i] = new List<int>();
        }
        for (int i = 1; i < n; ++i) {
            g[parents[i]].Add(i);
        }

        Dfs(0, -1);
        return ans;
    }

    private int Dfs(int i, int fa) {
        int cnt = 1;
        long score = 1;

        foreach (int j in g[i]) {
            if (j != fa) {
                int t = Dfs(j, i);
                cnt += t;
                score *= t;
            }
        }

        if (n - cnt > 0) {
            score *= n - cnt;
        }

        if (mx < score) {
            mx = score;
            ans = 1;
        } else if (mx == score) {
            ++ans;
        }

        return cnt;
    }
}

方法二：显式栈

思考

去掉结点后，分数是各连通块大小的乘积。 \(n\) 可以到 \(10^5\)。一条链上第一次递归总是走向唯一的孩子，深度就是 \(n\)，会超出默认递归上限。

分数要等每个孩子的子树大小都确定后才能相乘，结点上方剩余部分的大小是 \(n\) 减去这棵子树的大小。

因此用显式栈做后序。进入结点时先压入离开标记，再压入孩子。离开时把各孩子的子树大小相乘，若还有上方的块再乘 \(n-\textit{cnt}\)，并维护最高分及其出现次数。

我们先根据给定的父节点数组 parents 构建图 \(g\)，其中 \(g[i]\) 表示节点 \(i\) 的所有子节点。定义变量 \(\textit{ans}\) 表示最高得分的节点数目，变量 \(mx\) 表示最高得分。

用显式栈从根结点 \(0\) 做后序遍历。进入一个结点时，先压入该结点的离开标记，再压入它的孩子，这样孩子会先处理完。离开时，各孩子的子树大小已经写好。初始化 \(\textit{cnt} = 1\)、 \(\textit{score} = 1\)。把每个孩子的子树大小累乘到 \(\textit{score}\)，并累加到 \(\textit{cnt}\)。如果 \(n - \textit{cnt} > 0\)，再把 \(n - \textit{cnt}\) 累乘到 \(\textit{score}\)。若 \(\textit{score}\) 大于 \(mx\)，则把 \(mx\) 更新为 \(\textit{score}\)、 \(\textit{ans}\) 更新为 \(1\)；若相等，则把 \(\textit{ans}\) 加 \(1\)。最后把 \(\textit{cnt}\) 记为当前结点的子树大小。

时间复杂度 \(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
30
31
32
class Solution:
    def countHighestScoreNodes(self, parents: List[int]) -> int:
        n = len(parents)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[parents[i]].append(i)
        ans = mx = 0
        sz = [0] * n
        stk = [(0, -1, 0)]
        while stk:
            i, fa, state = stk.pop()
            if state == 0:
                stk.append((i, fa, 1))
                for j in g[i]:
                    if j != fa:
                        stk.append((j, i, 0))
            else:
                cnt = score = 1
                for j in g[i]:
                    if j != fa:
                        t = sz[j]
                        score *= t
                        cnt += t
                if n - cnt:
                    score *= n - cnt
                if mx < score:
                    mx = score
                    ans = 1
                elif mx == score:
                    ans += 1
                sz[i] = cnt
        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
38
39
40
41
42
43
44
45
46
47
48
class Solution {
    public int countHighestScoreNodes(int[] parents) {
        int n = parents.length;
        List<Integer>[] g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (int i = 1; i < n; ++i) {
            g[parents[i]].add(i);
        }
        int ans = 0;
        long mx = 0;
        int[] sz = new int[n];
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push(new int[] {i, fa, 1});
                for (int j : g[i]) {
                    if (j != fa) {
                        stk.push(new int[] {j, i, 0});
                    }
                }
            } else {
                int cnt = 1;
                long score = 1;
                for (int j : g[i]) {
                    if (j != fa) {
                        int t = sz[j];
                        cnt += t;
                        score *= t;
                    }
                }
                if (n - cnt > 0) {
                    score *= n - cnt;
                }
                if (mx < score) {
                    mx = score;
                    ans = 1;
                } else if (mx == score) {
                    ++ans;
                }
                sz[i] = cnt;
            }
        }
        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
38
39
40
41
42
43
44
45
46
47
48
class Solution {
public:
    int countHighestScoreNodes(vector<int>& parents) {
        int n = parents.size();
        vector<vector<int>> g(n);
        for (int i = 1; i < n; ++i) {
            g[parents[i]].push_back(i);
        }
        int ans = 0;
        long long mx = 0;
        vector<int> sz(n);
        vector<array<int, 3>> stk{{0, -1, 0}};
        while (!stk.empty()) {
            auto cur = stk.back();
            stk.pop_back();
            int i = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.push_back({i, fa, 1});
                for (int j : g[i]) {
                    if (j != fa) {
                        stk.push_back({j, i, 0});
                    }
                }
            } else {
                int cnt = 1;
                long long score = 1;
                for (int j : g[i]) {
                    if (j != fa) {
                        int t = sz[j];
                        cnt += t;
                        score *= t;
                    }
                }
                if (n - cnt) {
                    score *= n - cnt;
                }
                if (mx < score) {
                    mx = score;
                    ans = 1;
                } else if (mx == score) {
                    ++ans;
                }
                sz[i] = cnt;
            }
        }
        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
38
39
40
41
42
43
func countHighestScoreNodes(parents []int) (ans int) {
    n := len(parents)
    g := make([][]int, n)
    for i := 1; i < n; i++ {
        g[parents[i]] = append(g[parents[i]], i)
    }
    mx := 0
    sz := make([]int, n)
    stk := [][3]int{{0, -1, 0}}
    for len(stk) > 0 {
        cur := stk[len(stk)-1]
        stk = stk[:len(stk)-1]
        i, fa, state := cur[0], cur[1], cur[2]
        if state == 0 {
            stk = append(stk, [3]int{i, fa, 1})
            for _, j := range g[i] {
                if j != fa {
                    stk = append(stk, [3]int{j, i, 0})
                }
            }
        } else {
            cnt, score := 1, 1
            for _, j := range g[i] {
                if j != fa {
                    t := sz[j]
                    cnt += t
                    score *= t
                }
            }
            if n-cnt > 0 {
                score *= n - cnt
            }
            if mx < score {
                mx = score
                ans = 1
            } else if mx == score {
                ans++
            }
            sz[i] = cnt
        }
    }
    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
34
35
36
37
38
39
40
41
42
43
function countHighestScoreNodes(parents: number[]): number {
    const n = parents.length;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (let i = 1; i < n; i++) {
        g[parents[i]].push(i);
    }
    let ans = 0;
    let mx = 0;
    const sz: number[] = Array(n).fill(0);
    const stk: number[][] = [[0, -1, 0]];
    while (stk.length) {
        const [i, fa, state] = stk.pop()!;
        if (state === 0) {
            stk.push([i, fa, 1]);
            for (const j of g[i]) {
                if (j !== fa) {
                    stk.push([j, i, 0]);
                }
            }
        } else {
            let cnt = 1;
            let score = 1;
            for (const j of g[i]) {
                if (j !== fa) {
                    const t = sz[j];
                    cnt += t;
                    score *= t;
                }
            }
            if (n - cnt) {
                score *= n - cnt;
            }
            if (mx < score) {
                mx = score;
                ans = 1;
            } else if (mx === score) {
                ans++;
            }
            sz[i] = cnt;
        }
    }
    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
38
39
40
41
42
43
44
45
46
47
48
49
50
public class Solution {
    public int CountHighestScoreNodes(int[] parents) {
        int n = parents.Length;
        List<int>[] g = new List<int>[n];
        for (int i = 0; i < n; ++i) {
            g[i] = new List<int>();
        }
        for (int i = 1; i < n; ++i) {
            g[parents[i]].Add(i);
        }
        int ans = 0;
        long mx = 0;
        int[] sz = new int[n];
        Stack<int[]> stk = new Stack<int[]>();
        stk.Push(new int[] { 0, -1, 0 });
        while (stk.Count > 0) {
            int[] cur = stk.Pop();
            int i = cur[0], fa = cur[1], state = cur[2];
            if (state == 0) {
                stk.Push(new int[] { i, fa, 1 });
                foreach (int j in g[i]) {
                    if (j != fa) {
                        stk.Push(new int[] { j, i, 0 });
                    }
                }
            } else {
                int cnt = 1;
                long score = 1;
                foreach (int j in g[i]) {
                    if (j != fa) {
                        int t = sz[j];
                        cnt += t;
                        score *= t;
                    }
                }
                if (n - cnt > 0) {
                    score *= n - cnt;
                }
                if (mx < score) {
                    mx = score;
                    ans = 1;
                } else if (mx == score) {
                    ++ans;
                }
                sz[i] = cnt;
            }
        }
        return ans;
    }
}

评论