Polynomial Round 2022 E. Two Chess Pieces(dfs+dp)

E. Two Chess Pieces

题目大意：

给定n个节点的以1为根节点的有根树，现在在根节点上有两颗棋子，我们分别给他们规定了它们所必须经过的点，每次可以顺着树移动距离1，但是必须使得两颗棋子的距离不超过d，最终需要回到根节点，求最小的移动路径。

---

解题思路：

我们考虑棋子1,2他们必须经过的顶点，当棋子1必须经过顶点x时，我们根据距离约束可以确定，棋子2必须经过顶点x-d(与x距离为d的顶点)，不然就会超出距离约束。

考虑到这一点之后我们就可以先dfs一次记录所有顶点的x-d父节点到b数组中
之后我们只要将所有必须经过的点连起来就是需要通过的路径了
考虑f[2][N]表示棋子i是否必须经过顶点j，然后进行第二次dfs，自底向上更新f[i][j]即可。

最后考虑总路径贡献时：每当必须经过一个顶点时因为需要最终返回根节点，所以每个点的贡献为2，同时需要注意根节点本身不做贡献。

---

代码实现：

#include<bits/stdc++.h>
using namespace std;
# define int long long
# define endl "\n"
const int N = 2e5 + 10, inf = 0x3f3f3f3f;
int n,d;
vector<int> e[N];
int f[2][N];
int b[N];
int a[N];
void dfs1(int u,int fa,int dix){
	a[dix] = u;
	if(dix>d) b[u] = a[dix-d];//与点x距离为d的父节点 
	else b[u] = 1;
	
	for(auto v:e[u]){
		if(v == fa) continue;
		dfs1(v,u,dix+1);
	}
}
void dfs2(int u,int fa,int t){
	bool ok = false;
	for(auto v:e[u]){
		if(v == fa) continue;
		dfs2(v,u,t);
		ok |= f[t][v];//判断子节点是否需要经过 
	}
	f[t][u] |= ok;//如果子节点需要经过，父节点一定需要经过
	//自底向上更新 
}
void solve() {
	cin>>n>>d;
	for(int i = 1;i < n;++i){
		int a,b;
		cin>>a>>b;
		e[a].push_back(b);
		e[b].push_back(a);
	}
	dfs1(1,0,0);
	
	for(int i = 0;i < 2;++i){
		int m;
		cin>>m;
		for(int j = 1;j <= m;++j){
			int x;
			cin>>x;
			f[i][x] = f[i^1][b[x]] = 1;
			//如果棋子i需要经过点x,则另一个棋子需要经过点b[x] 
		}
	}
	
	dfs2(1,0,0);
	dfs2(1,0,1);
	
	int ans = 0;
	//计算贡献 
	for(int i = 2;i <= n;++i){
		ans += (f[0][i]==1)*2;
		ans += (f[1][i]==1)*2;
	}
	cout<<ans<<endl;
}
int tt;
signed main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	tt = 1;
//	cin >> tt;
	while (tt--) solve();
 
 
	return 0;
}