E. Two Chess Pieces -- (codeforces) 树形DP

发布时间 2023-07-11 16:32:37作者: N再再

原题链接:https://codeforces.com/contest/1774/problem/E

题意:两颗棋子,给出两颗棋子必须要去的顶点,且给出两颗棋子的相隔距离不能大于d,算出两颗棋子完成目标后走的距离。最后两颗棋子都要回到顶点1上。

思路:刚开始没想出来,顺着官方题解写的,大意就是我用数组 s1 和 s2 代表两颗棋子标记要去的顶点,最后算答案的时候直接计算就行了,标记要去的点主要是第一个 dfs1 就是用来标记两个棋子距离大于d后,会带动领一个棋子一定要走的路。

AC代码:

const int N = 1e6 +  10;
vector<int> v[N];
int s1[N], s2[N];
int a[N], b[N];
int n, d;

void dfs1(int u, int fa, int dis) // b数组用于标记当一个棋子走后至少另一个棋子也要在的点。
{
	a[dis] = u;
	if (dis > d) b[u] = a[dis - d];
	else b[u] = 1;
	
	for (int i = 0; i < v[u].size(); i ++)
	{
		int son = v[u][i];
		if (son == fa) continue;
		dfs1(son, u, dis + 1);
	}
}

void dfs2(int u, int fa) // 很模板的遍历走整棵树的路,当子节点有必须要走的点时,当前这个节点也要走
{
	bool flag = 0;
	for (int i = 0; i < v[u].size(); i ++)
	{
		int son = v[u][i];
		if (son == fa) continue;
		dfs2(son, u);
		flag |= s1[son]; // 判断子节点是否为要走的点
	}
	s1[u] |= flag;
}

void dfs3(int u, int fa) // 同理
{
	bool flag = 0;
	for (int i = 0; i < v[u].size(); i ++)
	{
		int son = v[u][i];
		if (son == fa) continue;
		dfs3(son, u);
		flag |= s2[son];
	}
	s2[u] |= flag;
}

void solve()
{
	cin >> n >> d;
	for (int i = 1; i < n; i ++)
	{
		int x, y;
		cin >> x >> y;
		v[x].push_back(y);
		v[y].push_back(x);
	}
	
	int m1, m2;
	dfs1(1, -1, 0);	
	cin >> m1;
	for (int i = 1; i <= m1; i ++) {
		int x;
		cin >> x, s1[x] = 1;	
		s2[b[x]] = 1;
	}
	
	cin >> m2;
	for (int i = 1; i <= m2; i ++) {
		int x;
		cin >> x, s2[x] = 1;
		s1[b[x]] = 1;	
	}
	dfs2(1, -1);
	dfs3(1, -1);
	int ans = 0;
	for (int i = 2; i <= n; i ++)
	{
		if (s1[i]) ans += 2;
		if (s2[i]) ans += 2;
	}
	cout << ans << '\n';
}