跳转至

3331. 修改后子树的大小

题目描述

给你一棵 n 个节点且根节点为编号 0 的树,节点编号为 0 到 n - 1 。这棵树用一个长度为 n 的数组 parent 表示,其中 parent[i] 是第 i 个节点的父亲节点的编号。由于节点 0 是根,parent[0] == -1 。

给你一个长度为 n 的字符串 s ,其中 s[i] 是节点 i 对应的字符。

对于节点编号从 1 到 n - 1 的每个节点 x ,我们 同时 执行以下操作 一次 :

  • 找到距离节点 x 最近 的祖先节点 y ,且 s[x] == s[y] 。
  • 如果节点 y 不存在,那么不做任何修改。
  • 否则,将节点 x 与它父亲节点之间的边 删除 ,在 x 与 y 之间连接一条边,使 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 <= 105
  • 对于所有的 i >= 1 ,都有 0 <= parent[i] <= n - 1 。
  • parent[0] == -1
  • parent 表示一棵合法的树。
  • s 只包含小写英文字母。

解法

方法一

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
    def findSubtreeSizes(self, parent: List[int], s: str) -> List[int]:
        def dfs(i: int, fa: int):
            ans[i] = 1
            d[s[i]].append(i)
            for j in g[i]:
                dfs(j, i)
            k = fa
            if len(d[s[i]]) > 1:
                k = d[s[i]][-2]
            if k != -1:
                ans[k] += ans[i]
            d[s[i]].pop()

        n = len(s)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[parent[i]].append(i)
        d = defaultdict(list)
        ans = [0] * n
        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
class Solution {
    private List<Integer>[] g;
    private List<Integer>[] d;
    private char[] s;
    private int[] ans;

    public int[] findSubtreeSizes(int[] parent, String s) {
        int n = s.length();
        g = new List[n];
        d = new List[26];
        this.s = s.toCharArray();
        Arrays.setAll(g, k -> new ArrayList<>());
        Arrays.setAll(d, k -> new ArrayList<>());
        for (int i = 1; i < n; ++i) {
            g[parent[i]].add(i);
        }
        ans = new int[n];
        dfs(0, -1);
        return ans;
    }

    private void dfs(int i, int fa) {
        ans[i] = 1;
        int idx = s[i] - 'a';
        d[idx].add(i);
        for (int j : g[i]) {
            dfs(j, i);
        }
        int k = d[idx].size() > 1 ? d[idx].get(d[idx].size() - 2) : fa;
        if (k >= 0) {
            ans[k] += ans[i];
        }
        d[idx].remove(d[idx].size() - 1);
    }
}
 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
class Solution {
public:
    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        int n = s.size();
        vector<int> g[n];
        vector<int> d[26];
        for (int i = 1; i < n; ++i) {
            g[parent[i]].push_back(i);
        }
        vector<int> ans(n);
        auto dfs = [&](this auto&& dfs, int i, int fa) -> void {
            ans[i] = 1;
            int idx = s[i] - 'a';
            d[idx].push_back(i);
            for (int j : g[i]) {
                dfs(j, i);
            }
            int k = d[idx].size() > 1 ? d[idx][d[idx].size() - 2] : fa;
            if (k >= 0) {
                ans[k] += ans[i];
            }
            d[idx].pop_back();
        };
        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
func findSubtreeSizes(parent []int, s string) []int {
    n := len(s)
    g := make([][]int, n)
    for i := 1; i < n; i++ {
        g[parent[i]] = append(g[parent[i]], i)
    }
    d := [26][]int{}
    ans := make([]int, n)
    var dfs func(int, int)
    dfs = func(i, fa int) {
        ans[i] = 1
        idx := int(s[i] - 'a')
        d[idx] = append(d[idx], i)
        for _, j := range g[i] {
            dfs(j, i)
        }
        k := fa
        if len(d[idx]) > 1 {
            k = d[idx][len(d[idx])-2]
        }
        if k != -1 {
            ans[k] += ans[i]
        }
        d[idx] = d[idx][:len(d[idx])-1]
    }
    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
function findSubtreeSizes(parent: number[], s: string): number[] {
    const n = parent.length;
    const g: number[][] = Array.from({ length: n }, () => []);
    const d: number[][] = Array.from({ length: 26 }, () => []);
    for (let i = 1; i < n; ++i) {
        g[parent[i]].push(i);
    }
    const ans: number[] = Array(n).fill(1);
    const dfs = (i: number, fa: number): void => {
        const idx = s.charCodeAt(i) - 97;
        d[idx].push(i);
        for (const j of g[i]) {
            dfs(j, i);
        }
        const k = d[idx].length > 1 ? d[idx].at(-2)! : fa;
        if (k >= 0) {
            ans[k] += ans[i];
        }
        d[idx].pop();
    };
    dfs(0, -1);
    return ans;
}

评论