子树中标签相同的节点数 图解题解

这道题到底在问什么

输入

n=7, 节点 0 与节点 2 都是标签 a

输出

ans[0]=2(整棵树里和根同标签 a 的有 2 个)

输入

任意一个叶子节点

输出

1(它的子树只有它自己)

输入

四个节点连成链、标签全是 b

输出

[4,3,2,1](越往上同标签越多)

最优解：一步一步想明白

1. 3记牢一句话: 每个节点向上交一份「子树里各字母出现几次」的统计, 父节点把孩子的统计合并再加上自己; 而 ans[i] 就是这份统计里节点 i 自己那个标签的次数。下面每一帧都在套这条规则。
2. 4根 0, 每个节点带一个字母标签先看清这棵树。紫色的节点 0 是根。它的两个孩子是 1 和 2; 节点 1 下面挂着 4 和 5, 节点 2 下面挂着 3 和 6。右边「标签表」记着每个节点的字母: 0 号 a, 1 号 b, 2 号 a, 3 号 e, 4 号 d, 5 号 c, 6 号 d。我们要给七个节点各算一个 ans, ans[i] 是它子树里和它同标签的节点数。
3. 5边是无向的, 从 0 当根定父子方向动手前先建邻接表。题目给的是无向边, 每条边的两个端点都要互相记成对方的邻居, 这样从任意点都能往四周走。然后把节点 0 当根, 从它出发往下走, 沿途遇到的节点就成了孩子。递归时带上父节点, 遇到父节点就跳过, 免得从孩子又走回父亲绕圈。方向定好, 这棵树就有了上下层级。
4. 6节点向上交一份 26 字母统计把规则在画面上钉死。每个节点要向上交一份统计: 它的子树里, 26 个小写字母分别出现了几次。父节点把所有孩子的统计逐字母相加, 再给自己的标签加 1, 就得到自己的统计。而某个节点的答案, 就是它这份统计里、自己那个标签出现的次数。带着这条规则, 开始后序搜索。
5. 7进入节点 0(标签 a), 结果待定后序深度优先搜索从根 0 开始。蓝色表示这个节点已经进入、正挂在搜索路径上, 但它的统计还没算出来, 因为得先把孩子那边都收完。节点 0 的标签是 a。先往它的第一个孩子 1 走下去。
6. 8进入节点 1(标签 b), 继续往下走到节点 1, 它的标签是 b。现在还不能下结论, 因为它下面还挂着 4 和 5 两个孩子, 得先把孩子的统计收上来才能合并。继续往它的孩子里下探, 先去 4。
7. 9节点 4(标签 d)是叶子走到节点 4, 标签是 d。它没有孩子, 是个叶子, 可以马上结算。叶子的子树就它自己一个节点, 统计也最简单。
8. 10叶子 4 统计 {d×1}, ans[4]=1结算节点 4。它变绿, 代表这棵子树算完了。子树里只有它自己, 标签是 d, 所以统计是 d 出现 1 次。节点 4 自己的标签就是 d, 在统计里 d 出现 1 次, 所以 ans[4] 等于 1。节点 4 把这份 {d×1} 的统计交回给父亲 1。接着看 1 的另一个孩子 5。
9. 11节点 5(标签 c)是叶子走到节点 5, 标签是 c, 它也是叶子, 同样可以直接结算。和刚才的 4 一样, 子树里只有它自己一个节点。
10. 12叶子 5 统计 {c×1}, ans[5]=1结算节点 5。它变绿, 子树里只有自己, 标签是 c, 统计就是 c 出现 1 次。节点 5 自己的标签 c 在统计里出现 1 次, 所以 ans[5] 等于 1。现在节点 1 的两个孩子都算完了: 4 交回 {d×1}, 5 交回 {c×1}。回到节点 1, 该把它们合并了。
11. 13合并孩子 {d×1} 与 {c×1}回到节点 1。把两个孩子的统计逐字母相加: 孩子 4 带回 d 出现 1 次, 孩子 5 带回 c 出现 1 次, 两个字母不重叠, 合并后是 c 一次、d 一次。还差最后一步, 把节点 1 自己的标签也加进去。
12. 14加自己 b, ans[1]=b 的次数=1结算节点 1。它自己的标签是 b, 给统计里 b 加 1, 子树 1 的完整统计就是 b 一次、c 一次、d 一次。节点 1 变绿。它自己的标签 b 在这份统计里出现 1 次, 所以 ans[1] 等于 1。注意子树 1 里虽然有三个节点, 但标签和 b 相同的只有节点 1 自己。节点 1 把这份统计交还给根 0。回根 0, 去看另一个孩子 2。
13. 15进入节点 2(标签 a), 先看孩子从根 0 走到另一个孩子 2, 它的标签是 a。它下面还挂着 3 和 6, 按后序还是要先把孩子看完, 才能合并出自己的统计。先去 3。
14. 16节点 3(标签 e)是叶子走到节点 3, 标签是 e, 叶子, 直接结算。子树里只有它自己一个节点。
15. 17叶子 3 统计 {e×1}, ans[3]=1结算节点 3。它变绿, 子树只有自己, 标签 e, 统计是 e 出现 1 次。自己的标签 e 出现 1 次, 所以 ans[3] 等于 1。节点 3 把 {e×1} 交回父亲 2。再看 2 的另一个孩子 6。
16. 18节点 6(标签 d)是叶子走到节点 6, 标签是 d, 也是叶子。和前面几个叶子一样, 子树里只有它自己。
17. 19叶子 6 统计 {d×1}, ans[6]=1结算节点 6。它变绿, 标签 d, 统计是 d 出现 1 次, ans[6] 等于 1。现在节点 2 的两个孩子都算完了: 3 交回 {e×1}, 6 交回 {d×1}。回到节点 2, 把它们合并。
18. 20合并孩子 {e×1} 与 {d×1}回到节点 2。把孩子的统计逐字母相加: 孩子 3 带回 e 一次, 孩子 6 带回 d 一次, 合并后是 d 一次、e 一次。下一步把节点 2 自己的标签 a 加进去。
19. 21加自己 a, ans[2]=a 的次数=1结算节点 2。它自己标签是 a, 给统计里 a 加 1, 子树 2 的完整统计是 a 一次、d 一次、e 一次。节点 2 变绿。它自己的标签 a 在统计里出现 1 次, 所以 ans[2] 等于 1。节点 2 把这份统计交还给根 0。现在根 0 的两个孩子都收齐了。
20. 22合并子树1 与子树2 的统计回到根 0。把两支孩子的统计逐字母相加。子树 1 带回 b 一次、c 一次、d 一次; 子树 2 带回 a 一次、d 一次、e 一次。两份相加: a 一次, b 一次, c 一次, d 在两边各一次合成 2 次, e 一次。最后还要把根自己的标签 a 加进去。
21. 23加自己 a, ans[0]=a 的次数=2结算根 0。它自己标签是 a, 给统计里 a 再加 1, a 变成 2 次。整棵树的统计是 a 两次、b 一次、c 一次、d 两次、e 一次。根 0 自己的标签是 a, 在统计里 a 出现 2 次, 所以 ans[0] 等于 2。这两个 a 来自节点 0 自己和节点 2。整棵树搜完, 答案出来了。
22. 24ans = [2,1,1,1,1,1,1]回放一遍。整个过程只做了一次后序遍历: 从根下探到每个叶子, 再自底向上把每层的统计合并回来, 七个节点的 ans 一次就全算出来了。结果是 ans 等于 2, 1, 1, 1, 1, 1, 1。除了根节点是 2, 其它六个都是 1, 因为它们各自的子树里没有第二个同标签的节点。
23. 25根统计 a×2 也有 d×2, ans 只看 a最后看一处最容易绕晕的地方。根 0 的统计里, 字母 a 出现 2 次, 字母 d 也出现 2 次(来自节点 4 和节点 6)。可 ans[0] 只等于 2, 是因为它只看根自己那个标签 a 的次数。d 出现几次和根没关系。这就是核心: 合并时统计全部 26 个字母, 但每个节点的答案只读出它自己那个标签的次数。

⚠️ 容易写错的地方

完整代码（Python / C++ / Java）

from __future__ import annotations
from typing import *
from collections import *
from functools import *
from itertools import *
from math import *
from heapq import *
from bisect import *
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next
class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
        def dfs(i, fa):
            ans[i] -= cnt[labels[i]]
            cnt[labels[i]] += 1
            for j in g[i]:
                if j != fa:
                    dfs(j, i)
            ans[i] += cnt[labels[i]]

        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        cnt = Counter()
        ans = [0] * n
        dfs(0, -1)
        return ans

#include <algorithm>
#include <array>
#include <cctype>
#include <climits>
#include <cmath>
#include <deque>
#include <functional>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <optional>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
using namespace std;
struct TreeNode {
  int val;
  TreeNode *left;
  TreeNode *right;
  TreeNode() : val(0), left(nullptr), right(nullptr) {}
  TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
  TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};
struct ListNode {
  int val;
  ListNode *next;
  ListNode() : val(0), next(nullptr) {}
  ListNode(int x) : val(x), next(nullptr) {}
  ListNode(int x, ListNode *next) : val(x), next(next) {}
};
class Solution {
public:
    vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
        vector<vector<int>> g(n);
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        vector<int> ans(n);
        int cnt[26]{};
        function<void(int, int)> dfs = [&](int i, int fa) {
            int k = labels[i] - 'a';
            ans[i] -= cnt[k];
            cnt[k]++;
            for (int& j : g[i]) {
                if (j != fa) {
                    dfs(j, i);
                }
            }
            ans[i] += cnt[k];
        };
        dfs(0, -1);
        return ans;
    }
};

import java.util.*;
class TreeNode { int val; TreeNode left; TreeNode right; TreeNode(){} TreeNode(int val){this.val=val;} TreeNode(int val, TreeNode left, TreeNode right){this.val=val;this.left=left;this.right=right;} }
class ListNode { int val; ListNode next; ListNode(){} ListNode(int val){this.val=val;} ListNode(int val, ListNode next){this.val=val;this.next=next;} }
class Solution {
    private List<Integer>[] g;
    private String labels;
    private int[] ans;
    private int[] cnt;

    public int[] countSubTrees(int n, int[][] edges, String labels) {
        g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (int[] e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        this.labels = labels;
        ans = new int[n];
        cnt = new int[26];
        dfs(0, -1);
        return ans;
    }

    private void dfs(int i, int fa) {
        int k = labels.charAt(i) - 'a';
        ans[i] -= cnt[k];
        cnt[k]++;
        for (int j : g[i]) {
            if (j != fa) {
                dfs(j, i);
            }
        }
        ans[i] += cnt[k];
    }
}

复杂度

面试官可能追问