2025-05-20：修改后子树的大小。用go语言，给你一棵有 n 个节点的树，节点编号从 0 到 n-1，且节点 0 是树的根

2025-05-20：修改后子树的大小。用go语言，给你一棵有 n 个节点的树，节点编号从 0 到 n-1，且节点 0 是树的根。树的结构用一个长度为 n 的数组 parent 表示，其中 parent[i] 表示节点 i 的父节点编号。由于节点 0 是根节点，所以 parent[0] = -1。

同时给你一个长度为 n 的字符串 s，s[i] 表示节点 i 对应的字符。

对每个编号为 1 到 n-1 的节点 x，执行如下操作一次：

1. 1. 找出节点 x 最近的祖先节点 y，且满足 s[x] == s[y]。
2. 2. 如果找不到这样的祖先节点 y，则不做任何操作。
3. 3. 如果找到了，则断开节点 x 和它原父节点之间的连接，改为将边连接到节点 y，使 y 成为 x 新的父节点。

执行完这些操作后，返回一个长度为 n 的数组 answer，其中 answer[i] 表示最终树中，以节点 i 为根的子树包含的节点个数。

n == parent.length == s.length。

1 <= n <= 100000。

对于所有的 i >= 1 ，都有 0 <= parent[i] <= n - 1 。

parent[0] == -1。

parent 表示一棵合法的树。

s 只包含小写英文字母。

输入：parent = [-1,0,0,1,1,1], s = "abaabc"。

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

解释：

在这里插入图片描述

节点 3 的父节点从节点 1 变为节点 0 。

题目来自力扣3331。

初始理解题目

我们有一棵树，节点编号从0到n-1，其中0是根节点。树的连接关系由parent数组给出，parent[i]表示节点i的父节点。同时，每个节点有一个对应的字符s[i]。我们需要对每个非根节点（即节点1到n-1）执行一次操作：

1. 1. 找到节点x的最近的祖先y，使得s[x] == s[y]。
2. 2. 如果找到这样的y，就将x的父节点从原来的parent[x]改为y。
3. 3. 如果没有这样的y，就不做任何操作。

执行完所有操作后，我们需要计算每个节点作为根的子树的大小（即子树中包含的节点数量）。

执行操作

我们需要对节点1到5分别执行操作：

1. 节点1：

• • 字符是'b'。
• • 祖先：0（字符'a'）。
• • 没有与's[1]'相同的祖先（'a' != 'b'），所以不操作。

2. 节点2：

• • 字符是'a'。
• • 祖先：0（字符'a'）。
• • 找到最近的相同字符的祖先是0。
• • 将节点2的父节点从0改为0（实际上没有变化），所以不操作。

3. 节点3：

• • 字符是'a'。
• • 祖先路径：1（'b'） -> 0（'a'）。
• • 最近的相同字符的祖先是0（'a'）。
• • 将节点3的父节点从1改为0。
• • 修改后的树： 0 (a) / | \ 1(b)2(a)3(a) / \ 4(b)5(c)

4. 节点4：

• • 字符是'b'。
• • 祖先路径：1（'b'） -> 0（'a'）。
• • 最近的相同字符的祖先是1（'b'）。
• • 将节点4的父节点从1改为1（没有变化），所以不操作。

5. 节点5：

• • 字符是'c'。
• • 祖先路径：1（'b'） -> 0（'a'）。
• • 没有与's[5]'相同的祖先（'b' != 'c'，'a' != 'c'），所以不操作。

最终树的结构：

        0 (a)
      / | \
    1(b)2(a)3(a)
      / \
    4(b)5(c)

计算子树大小

现在我们需要计算每个节点作为根的子树的大小：

• • 节点5：子树{5}，大小1。
• • 节点4：子树{4}，大小1。
• • 节点3：子树{3}，大小1。
• • 节点2：子树{2}，大小1。
• • 节点1：子树{1, 4, 5}，大小3。
• • 节点0：子树{0, 1, 2, 3, 4, 5}，大小6。

因此，answer = [6, 3, 1, 1, 1, 1]。

代码逻辑分析

代码的主要逻辑是通过深度优先搜索（DFS）来计算子树大小。具体步骤如下：

1. 1. 构建树结构：根据parent数组构建邻接表g，表示每个节点的子节点。
2. 2. DFS遍历：
   • • 维护一个size数组，表示每个节点的子树大小。
   • • 维护一个ancestor数组（长度为26，对应小写字母），记录当前路径中每个字符最近的祖先节点。
   • • 对于当前节点x：
     • • 设置size[x] = 1（至少包含自己）。
     • • 记录当前字符s[x]的旧祖先，并将ancestor[s[x]]更新为x。
     • • 递归处理所有子节点y：
       • • 对于子节点y，找到其字符s[y]的最近祖先anc（如果ancestor[s[y]]为-1，则用x作为默认祖先）。
       • • 将size[anc]增加size[y]（因为y的子树现在属于anc的子树）。
     • • 恢复ancestor[s[x]]为旧值（回溯）。
3. 3. 结果：size数组即为所求的answer。

时间复杂度

• • 构建邻接表：O(n)。
• • DFS遍历：每个节点被访问一次，每次访问处理其子节点和ancestor数组的操作是O(1)（因为字母表是固定大小的26）。
• • 总时间复杂度：O(n)。

空间复杂度

• • 邻接表g：O(n)。
• • size数组：O(n)。
• • ancestor数组：O(26) = O(1)。
• • DFS递归栈：最坏情况下是O(n)（树退化为链表时）。
• • 总空间复杂度：O(n)。

总结

• • 大体过程：
  1. 1. 构建树的邻接表。
  2. 2. 通过DFS遍历树，维护当前路径中每个字符的最近祖先。
  3. 3. 对于每个子节点，根据其字符找到最近的祖先，并将子树大小累加到该祖先。
  4. 4. 回溯时恢复祖先状态。
  5. 5. 最终size数组即为答案。
• • 时间复杂度：O(n)。
• • 空间复杂度：O(n)。

Go完整代码如下：

package main

import (
    "fmt"
)

func findSubtreeSizes(parent []int, s string) []int {
    n := len(parent)
    g := make([][]int, n)
    for i := 1; i < n; i++ {
        p := parent[i]
        g[p] = append(g[p], i)
    }

    size := make([]int, n)
    ancestor := [26]int{}
    for i := range ancestor {
        ancestor[i] = -1
    }
    var dfs func(int)
    dfs = func(x int) {
        size[x] = 1
        sx := s[x] - 'a'
        old := ancestor[sx]
        ancestor[sx] = x
        for _, y := range g[x] {
            dfs(y)
            anc := ancestor[s[y]-'a']
            if anc < 0 {
                anc = x
            }
            size[anc] += size[y]
        }
        ancestor[sx] = old // 恢复现场
    }
    dfs(0)
    return size
}

func main() {
    parent := []int{-1, 0, 0, 1, 1, 1}
    s := "abaabc"
    result := findSubtreeSizes(parent, s)
    fmt.Println(result)
}

Rust完整代码如下：

fn find_subtree_sizes(parent: Vec<i32>, s: &str) ->Vec<i32> {
    letn = parent.len();
    letmut g = vec![vec![]; n];
    foriin1..n {
        letp = parent[i] asusize;
        g[p].push(i);
    }

    letmut size = vec![0; n];
    // 记录字符对应最近祖先节点编号，-1 表示无
    letmut ancestor = vec![-1; 26];

    lets_bytes = s.as_bytes();

    fndfs(
        x: usize,
        g: &Vec<Vec<usize>>,
        s_bytes: &[u8],
        ancestor: &mutVec<i32>,
        size: &mutVec<i32>,
    ) {
        size[x] = 1;
        letsx = (s_bytes[x] - b'a') asusize;
        letold = ancestor[sx];
        ancestor[sx] = x asi32;

        for &y in &g[x] {
            dfs(y, g, s_bytes, ancestor, size);

            letanc = ancestor[(s_bytes[y] - b'a') asusize];
            letanc_index = if anc < 0 { x asi32 } else { anc };

            size[anc_index asusize] += size[y];
        }

        ancestor[sx] = old; // 恢复现场
    }

    dfs(0, &g, s_bytes, &mut ancestor, &mut size);
    size
}

fnmain() {
    letparent = vec![-1, 0, 0, 1, 1, 1];
    lets = "abaabc";

    letresult = find_subtree_sizes(parent, s);
    println!("{:?}", result);
}