codeforces-go/leetcode/biweekly/142/b at master · EndlessCheng/codeforces-go

Name	Name	Last commit message	Last commit date
parent directory ..
README.md	README.md
b.go	b.go
b.txt	b.txt
b_test.go	b_test.go

方法一：两次 DFS

第一次 DFS

写一个自顶向下的 DFS，同时维护路径上的每个字母对应的最深节点 $ancestor$ ，初始化成一个全为 $- 1$ 的数组（或者哈希表）。递归到节点 $x$ 时，更新 $ancestor [s [x]] = x$ 。

例如示例 1，递归到节点 $1$ 时，有 $ancestor [a] = 0, ancestor [b] = 1$ 。

对于节点 $x$ 及其儿子 $y$ ，如果 $ancestor [s [y]] \neq - 1$ ，则把 $y$ 加到 $ancestor [s [y]]$ 的子节点列表中，同时把 $x$ 的儿子 $y$ 改成 $- 1$ ，在第二次 DFS 中，不去递归等于 $- 1$ 的儿子。

注意要写回溯，递归返回前，要恢复 $ancestor [s [x]]$ 的原始值（递归到 $x$ 之前的值）。

第二次 DFS

写一个自底向上的 DFS，返回子树大小。

当前子树大小为 $1$ （自己）加上所有儿子子树的大小之和。

⚠注意：在遍历列表的同时，往列表中加入数据，会导致我们遍历到新加入列表的数据。可以考虑固定遍历次数为 $x$ 的子节点列表的长度。

Go 语言不用考虑这个问题。

class Solution:
    def findSubtreeSizes(self, parent: List[int], s: str) -> List[int]:
        n = len(parent)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[parent[i]].append(i)

        ancestor = defaultdict(lambda: -1)
        def rebuild(x: int) -> None:
            old = ancestor[s[x]]
            ancestor[s[x]] = x
            for i in range(len(g[x])):
                y = g[x][i]
                if (anc := ancestor[s[y]]) != -1:
                    g[anc].append(y)
                    g[x][i] = -1  # -1 表示删除 y
                rebuild(y)
            ancestor[s[x]] = old  # 恢复现场
        rebuild(0)

        size = [1] * n  # 注意这里已经把 1 算进去了
        def dfs(x: int) -> None:
            for y in g[x]:
                if y != -1:  # y 没被删除
                    dfs(y)
                    size[x] += size[y]
        dfs(0)
        return size

class Solution {
    public int[] findSubtreeSizes(int[] parent, String s) {
        int n = parent.length;
        List<Integer>[] g = new ArrayList[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int i = 1; i < n; i++) {
            g[parent[i]].add(i);
        }

        int[] ancestor = new int[26];
        Arrays.fill(ancestor, -1);
        rebuild(0, g, s.toCharArray(), ancestor);

        int[] size = new int[n];
        dfs(0, g, size);
        return size;
    }

    private void rebuild(int x, List<Integer>[] g, char[] s, int[] ancestor) {
        int sx = s[x] - 'a';
        int old = ancestor[sx];
        ancestor[sx] = x;
        for (int i = g[x].size() - 1; i >= 0; i--) {
            int y = g[x].get(i);
            int anc = ancestor[s[y] - 'a'];
            if (anc != -1) {
                g[anc].add(y);
                g[x].set(i, -1); // -1 表示删除 y
            }
            rebuild(y, g, s, ancestor);
        }
        ancestor[sx] = old; // 恢复现场
    }

    private void dfs(int x, List<Integer>[] g, int[] size) {
        size[x] = 1;
        for (int y : g[x]) {
            if (y != -1) { // y 没被删除
                dfs(y, g, size);
                size[x] += size[y];
            }
        }
    }
}

class Solution {
public:
    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        int n = parent.size();
        vector<vector<int>> g(n);
        for (int i = 1; i < n; i++) {
            g[parent[i]].push_back(i);
        }

        int ancestor[26];
        ranges::fill(ancestor, -1);
        auto rebuild = [&](auto&& rebuild, int x) -> void {
            int sx = s[x] - 'a';
            int old = ancestor[sx];
            ancestor[sx] = x;
            for (int i = g[x].size() - 1; i >= 0; i--) {
                int y = g[x][i];
                int anc = ancestor[s[y] - 'a'];
                if (anc != -1) {
                    g[anc].push_back(y);
                    g[x][i] = -1; // -1 表示删除 y
                }
                rebuild(rebuild, y);
            }
            ancestor[sx] = old; // 恢复现场
        };
        rebuild(rebuild, 0);

        vector<int> size(n, 1); // 注意这里已经把 1 算进去了
        auto dfs = [&](auto&& dfs, int x) -> void {
            for (int y : g[x]) {
                if (y != -1) { // y 没被删除
                    dfs(dfs, y);
                    size[x] += size[y];
                }
            }
        };
        dfs(dfs, 0);
        return size;
    }
};

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)
	}

	ancestor := [26]int{}
	for i := range ancestor {
		ancestor[i] = -1
	}
	var rebuild func(int)
	rebuild = func(x int) {
		sx := s[x] - 'a'
		old := ancestor[sx]
		ancestor[sx] = x
		for i, y := range g[x] {
			if anc := ancestor[s[y]-'a']; anc != -1 {
				g[anc] = append(g[anc], y)
				g[x][i] = -1 // -1 表示删除 y
			}
			rebuild(y)
		}
		ancestor[sx] = old // 恢复现场
	}
	rebuild(0)

	size := make([]int, n)
	var dfs func(int)
	dfs = func(x int) {
		size[x] = 1
		for _, y := range g[x] {
			if y != -1 { // y 没被删除
				dfs(y)
				size[x] += size[y]
			}
		}
	}
	dfs(0)
	return size
}

复杂度分析

时间复杂度：$\mathcal{O}(n+|\Sigma|)$，其中 $n$ 是 $s$ 的长度，$|\Sigma|=26$ 是字符集合的大小。
空间复杂度：$\mathcal{O}(n)$。

方法二：一次 DFS

把两个 DFS 结合起来。在 $dfs (y)$ 结束后，分类讨论：

如果 $y$ 没有对应的祖先，把 $size [x]$ 增加 $size [y]$ 。
如果 $y$ 有对应的祖先 $anc$ ，把 $size [anc]$ 增加 $size [y]$ 。

具体请看视频讲解，欢迎点赞关注~

class Solution:
    def findSubtreeSizes(self, parent: List[int], s: str) -> List[int]:
        n = len(parent)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[parent[i]].append(i)

        size = [1] * n
        ancestor = defaultdict(lambda: -1)
        def dfs(x: int) -> None:
            old = ancestor[s[x]]
            ancestor[s[x]] = x
            for y in g[x]:
                dfs(y)
                anc = ancestor[s[y]]
                size[x if anc < 0 else anc] += size[y]
            ancestor[s[x]] = old  # 恢复现场
        dfs(0)
        return size

class Solution {
    public int[] findSubtreeSizes(int[] parent, String s) {
        int n = parent.length;
        List<Integer>[] g = new ArrayList[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int i = 1; i < n; i++) {
            g[parent[i]].add(i);
        }

        int[] size = new int[n];
        int[] ancestor = new int[26];
        Arrays.fill(ancestor, -1);
        dfs(0, g, s.toCharArray(), ancestor, size);
        return size;
    }

    private void dfs(int x, List<Integer>[] g, char[] s, int[] ancestor, int[] size) {
        size[x] = 1;
        int sx = s[x] - 'a';
        int old = ancestor[sx];
        ancestor[sx] = x;
        for (int y : g[x]) {
            dfs(y, g, s, ancestor, size);
            int anc = ancestor[s[y] - 'a'];
            size[anc < 0 ? x : anc] += size[y];
        }
        ancestor[sx] = old; // 恢复现场
    }
}

class Solution {
public:
    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        int n = parent.size();
        vector<vector<int>> g(n);
        for (int i = 1; i < n; i++) {
            g[parent[i]].push_back(i);
        }

        vector<int> size(n, 1);
        int ancestor[26];
        ranges::fill(ancestor, -1);
        auto dfs = [&](auto&& dfs, int x) -> void {
            int sx = s[x] - 'a';
            int old = ancestor[sx];
            ancestor[sx] = x;
            for (int y : g[x]) {
                dfs(dfs, y);
                int anc = ancestor[s[y] - 'a'];
                size[anc < 0 ? x : anc] += size[y];
            }
            ancestor[sx] = old; // 恢复现场
        };
        dfs(dfs, 0);
        return size;
    }
};

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
}

复杂度分析

时间复杂度：$\mathcal{O}(n+|\Sigma|)$，其中 $n$ 是 $s$ 的长度，$|\Sigma|=26$ 是字符集合的大小。
空间复杂度：$\mathcal{O}(n)$。

更多相似题目，见链表、二叉树与一般树中的「§2.2 自顶向下 DFS」和「§2.3 自底向上 DFS」。

分类题单

如何科学刷题？

我的题解精选（已分类）

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

b

b

README.md

方法一：两次 DFS

第一次 DFS

第二次 DFS

复杂度分析

方法二：一次 DFS

复杂度分析

分类题单

Files

b

Directory actions

More options

Directory actions

More options

Latest commit

History

b

Folders and files

parent directory

README.md

方法一：两次 DFS

第一次 DFS

第二次 DFS

复杂度分析

方法二：一次 DFS

复杂度分析

分类题单