3331. 找出变化后的子树大小

Medium

题目描述

给你一棵以节点 0 为根的树，由 n 个节点组成，节点编号从 0 到 n - 1。树用长度为 n 的数组 parent 表示，其中 parent[i] 是节点 i 的父节点。由于节点 0 是根，所以 parent[0] == -1。

同时给你一个长度为 n 的字符串 s，其中 s[i] 是分配给节点 i 的字符。

我们对树进行一次性的以下更改，对所有节点 x（从 1 到 n - 1）同时进行：

找到最近的节点 y，使得 y 是 x 的祖先，且 s[x] == s[y]。
如果节点 y 不存在，则不做任何操作。
否则，删除 x 与其当前父节点之间的边，并通过在它们之间添加边来使节点 y 成为 x 的新父节点。

返回一个长度为 n 的数组 answer，其中 answer[i] 是最终树中以节点 i 为根的子树的大小。

示例 1：

输入：parent = [-1,0,0,1,1,1], s = "abaabc"
输出：[6,3,1,1,1,1]
解释：节点 3 的父节点将从节点 1 变为节点 0。

示例 2：

输入：parent = [-1,0,4,0,1], s = "abbba"
输出：[5,2,1,1,1]
解释：同时发生以下变化：
- 节点 4 的父节点将从节点 1 变为节点 0。
- 节点 2 的父节点将从节点 4 变为节点 1。

约束条件：

n == parent.length == s.length
1 <= n <= 10⁵
对于所有 i >= 1，0 <= parent[i] <= n - 1
parent[0] == -1
parent 表示一棵有效的树
s 仅由小写英文字母组成

解题思路

这道题需要模拟树的重构过程，关键思路如下：

建立原始树结构：根据 parent 数组构建邻接表表示的原始树。
DFS 遍历并重新分配父节点：从根节点 0 开始进行深度优先搜索。在遍历过程中，维护一个栈来记录从根到当前节点路径上每个字符最近出现的节点位置。
查找新父节点：对于每个非根节点，在当前路径上寻找具有相同字符的最近祖先节点。如果找到，则该节点成为当前节点的新父节点；否则保持原父节点不变。
构建新树并计算子树大小：根据重新分配的父子关系构建新的树结构，然后通过 DFS 计算每个节点的子树大小。

核心技巧：

使用字符到节点的映射栈，在 DFS 过程中动态维护路径上每个字符的最新位置
进入节点时保存当前字符的旧映射，退出时恢复，确保回溯正确
分两阶段：先重构树结构，再计算子树大小

时间复杂度为 O(n)，空间复杂度为 O(n)，是最优解法。

代码实现

class Solution {
public:
    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        int n = parent.size();
        vector<vector<int>> children(n);
        
        // Build original tree
        for (int i = 1; i < n; i++) {
            children[parent[i]].push_back(i);
        }
        
        vector<int> newParent = parent;
        vector<int> charToNode(26, -1);
        
        function<void(int)> dfs = [&](int node) {
            int oldNode = charToNode[s[node] - 'a'];
            charToNode[s[node] - 'a'] = node;
            
            for (int child : children[node]) {
                int targetChar = s[child] - 'a';
                if (charToNode[targetChar] != -1) {
                    newParent[child] = charToNode[targetChar];
                }
                dfs(child);
            }
            
            charToNode[s[node] - 'a'] = oldNode;
        };
        
        dfs(0);
        
        // Build new tree and calculate subtree sizes
        vector<vector<int>> newChildren(n);
        for (int i = 1; i < n; i++) {
            newChildren[newParent[i]].push_back(i);
        }
        
        vector<int> result(n);
        function<int(int)> calcSize = [&](int node) -> int {
            int size = 1;
            for (int child : newChildren[node]) {
                size += calcSize(child);
            }
            return result[node] = size;
        };
        
        calcSize(0);
        return result;
    }
};

class Solution:
    def findSubtreeSizes(self, parent: List[int], s: str) -> List[int]:
        n = len(parent)
        children = [[] for _ in range(n)]
        
        # Build original tree
        for i in range(1, n):
            children[parent[i]].append(i)
        
        new_parent = parent[:]
        char_to_node = [-1] * 26
        
        def dfs(node):
            old_node = char_to_node[ord(s[node]) - ord('a')]
            char_to_node[ord(s[node]) - ord('a')] = node
            
            for child in children[node]:
                target_char = ord(s[child]) - ord('a')
                if char_to_node[target_char] != -1:
                    new_parent[child] = char_to_node[target_char]
                dfs(child)
            
            char_to_node[ord(s[node]) - ord('a')] = old_node
        
        dfs(0)
        
        # Build new tree and calculate subtree sizes
        new_children = [[] for _ in range(n)]
        for i in range(1, n):
            new_children[new_parent[i]].append(i)
        
        result = [0] * n
        
        def calc_size(node):
            size = 1
            for child in new_children[node]:
                size += calc_size(child)
            result[node] = size
            return size
        
        calc_size(0)
        return result

public class Solution {
    public int[] FindSubtreeSizes(int[] parent, string s) {
        int n = parent.Length;
        List<int>[] children = new List<int>[n];
        for (int i = 0; i < n; i++) {
            children[i] = new List<int>();
        }
        
        // Build original tree
        for (int i = 1; i < n; i++) {
            children[parent[i]].Add(i);
        }
        
        int[] newParent = new int[n];
        Array.Copy(parent, newParent, n);
        int[] charToNode = new int[26];
        Array.Fill(charToNode, -1);
        
        void Dfs(int node) {
            int oldNode = charToNode[s[node] - 'a'];
            charToNode[s[node] - 'a'] = node;
            
            foreach (int child in children[node]) {
                int targetChar = s[child] - 'a';
                if (charToNode[targetChar] != -1) {
                    newParent[child] = charToNode[targetChar];
                }
                Dfs(child);
            }
            
            charToNode[s[node] - 'a'] = oldNode;
        }
        
        Dfs(0);
        
        // Build new tree and calculate subtree sizes
        List<int>[] newChildren = new List<int>[n];
        for (int i = 0; i < n; i++) {
            newChildren[i] = new List<int>();
        }
        for (int i = 1; i < n; i++) {
            newChildren[newParent[i]].Add(i);
        }
        
        int[] result = new int[n];
        
        int CalcSize(int node) {
            int size = 1;
            foreach (int child in newChildren[node]) {
                size += CalcSize(child);
            }
            return result[node] = size;
        }
        
        CalcSize(0);
        return result;
    }
}

var findSubtreeSizes = function(parent, s) {
    const n = parent.length;
    const children = Array.from({length: n}, () => []);
    
    // Build original tree
    for (let i = 1; i < n; i++) {
        children[parent[i]].push(i);
    }
    
    const newParent = [...parent];
    const charToNode = new Array(26).fill(-1);
    
    function dfs(node) {
        const oldNode = charToNode[s.charCodeAt(node) - 97];
        charToNode[s.charCodeAt(node) - 97] = node;
        
        for (const child of children[node]) {
            const targetChar = s.charCodeAt(child) - 97;
            if (charToNode[targetChar] !== -1) {
                newParent[child] = charToNode[targetChar];
            }
            dfs(child);
        }
        
        charToNode[s.charCodeAt(node) - 97] = oldNode;
    }
    
    dfs(0);
    
    // Build new tree and calculate subtree sizes
    const newChildren = Array.from({length: n}, () => []);
    for (let i = 1; i < n; i++) {
        newChildren[newParent[i]].push(i);
    }
    
    const result = new Array(n);
    
    function calcSize(node) {
        let size = 1;
        for (const child of newChildren[node]) {
            size += calcSize(child);
        }
        return result[node] = size;
    }
    
    calcSize(0);
    return result;
};

在 LeetCode 上运行：力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型	值	说明
时间复杂度	O(n)	需要两次DFS遍历，每次访问所有节点一次
空间复杂度	O(n)	需要存储树的邻接表、新父节点数组和递归调用栈

题目描述#

解题思路#

代码实现#

复杂度分析#

题目描述

解题思路

代码实现

复杂度分析