力扣每日一题 统计可能的树根数目 换根DP

Problem: 2581. 统计可能的树根数目

思路

👨‍🏫 灵神

复杂度

时间复杂度: $O(n+m)$

空间复杂度: $O(n+m)$

Code

class Solution {
	List<Integer>[] g;// 邻接表
	Set<Long> set = new HashSet<>();// 猜测的父子关系
	int k, ans, cnt0;

	public int rootCount(int[][] edges, int[][] guesses, int k)
	{
		this.k = k;
		g = new ArrayList[edges.length + 1];
		Arrays.setAll(g, x -> new ArrayList<>());
		for (int[] e : edges)
		{
//			邻接表建图
			int x = e[0];
			int y = e[1];
			g[x].add(y);
			g[y].add(x);
		}
		for (int[] e : guesses)// 猜测转为 哈希表
			set.add((long) e[0] << 32 | e[1]);// 父<<32 | 子
		dfs(0, -1);
		reroot(0, -1, cnt0);// 计算以 0 为根时，猜对的个数 cnt0
		return ans;
	}

	/**
	 * @param x   当前根节点
	 * @param fa  当前根节点的
	 * @param cnt
	 */
	private void reroot(int x, int fa, int cnt)
	{
		if (cnt >= k)// 说明以当前节点为根 符合题目条件
			ans++;
		for (int y : g[x])// 找当前节点的下一个节点
		{
			if (y != fa)// 去重
			{
				int c = cnt;// 换根后重新计算猜对的数量（由 x->y 变成 y->x）
				if (set.contains((long) x << 32 | y))//对变错
					c--;
				if (set.contains((long) y << 32 | x))//错变对
					c++;
				reroot(y, x, c);//向下递归换根操作
			}
		}
	}

	/**
	 * @param x  当前节点
	 * @param fa 当前节点的父结点
	 */
	private void dfs(int x, int fa)
	{
		for (int y : g[x])// 遍历 x 的子结点
		{
			if (y != fa)
			{
				if (set.contains((long) x << 32 | y))// 猜对了
					cnt0++;
				dfs(y, x);// 递归遍历所有子结点
			}
		}
	}
}