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

每个节点要数出子树里和自己同标签的有几个

给一棵 n 个节点的树，编号 0 到 n-1，根是 0，每个节点带一个小写字母标签，记在 labels 里。一个节点的子树，就是它自己加上底下所有后代。要对每个 i 算出它子树里和它同标签的节点有几个，i 自己也算。题面 n=7，edges=[[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]]，labels=abaedcd，答案 [2,1,1,1,1,1,1]——只有根 0 是 2，因为标签 a 的有节点 0 和 2。

对每个节点各数一遍子树，长链上会退化

对每个节点都把整棵子树从头数一遍，就得重走一遍它的后代。碰上从根拉到底的长链越靠根子树越大：根数 n 个、下一个数 n-1 个，加起来是平方级 O(n²)，n 一大就走不动。省在孩子的子树早数过，父节点接过来直接用，不必重数。

先把无向边建成邻接表，再用后序把孩子收上来

题目给的边不分方向，得先建邻接表，也就是每个节点记一份自己连到谁的名单，每条边两个端点互相记成邻居。从根 0 往下走，走到谁谁就是孩子。因为两头都记了，递归要带上父节点、遇到就跳过，否则会在父子间来回无限走、撑爆栈。遍历用后序，也就是先把一个节点的孩子都处理完再处理自己——父节点合并时，孩子的统计已经摆好。

共享一个计数器，进节点前减、出子树后加

顺着这个思路，不给每个节点单独存统计，而是全程共用一个 26 格的计数器 cnt，记各字母眼下出现了几次。诀窍在进入节点 i、还没记数之前：先用 ans[i] 减去此刻 cnt 里它标签的值，把外面已有的同标签底数扣掉；再给自己标签加 1，递归所有孩子——孩子会把各自子树里的同标签也累进 cnt。等子树走完，让 ans[i] 加上此刻该标签的值。一减一加，净增的正好是 i 子树贡献的同标签数，i 自己也在其中。

拿题面这棵 7 节点的树，把差值法走通一遍

跟着差值法走一遍。进根 0（a），cnt[a] 从 0 到 1，ans[0] 进门减 0。进节点 1（b），孩子 4、5 是叶子，各自把 cnt[d]、cnt[c] 从 0 加到 1，ans[4]=ans[5]=1；回节点 1，ans[1]=1-0=1。

进节点 2（a），此刻 cnt[a] 已是 1、根记过，ans[2] 先减 1。孩子 3（e）给 ans[3]=1；孩子 6（d）进去前 cnt[d] 已是 1、节点 4 记的，出来变 2，ans[6]=2-1=1，差值只算 6 自己。出节点 2 时 cnt[a]=2，ans[2]=1。回根 0，ans[0]=2；注意 cnt 里 d 也是 2，但 ans[0] 只读 a，答案是 [2,1,1,1,1,1,1]。

为什么是 O(n)，三种极端小树各长什么样

建邻接表把 n-1 条边各记两遍是 O(n)，后序 DFS 每个节点只进出一次、计数器做常数次加减，也是 O(n)，合起来时间 O(n)。空间上邻接表 O(n)、递归栈最深是树高即长链时 O(n)，计数器不计，空间 O(n)。三种极端小树验一下：单节点答案是 1；一条链标签全同，每个答案等于它子树的大小；标签全不同则答案全是 1。

Python

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

C++

#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;
    }
};

Java

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];
    }
}