Codeforces 1363E - Tree Shuffling （思维/dfs/记忆化）

Posted 2020-12-23 stelayuri

Codeforces Round #646 (Div. 2) - E. Tree Shuffling

题意

有一棵 n 个节点 n-1 条边的树，以节点 1 为根节点

每个节点 i 拥有价值 a_i 和一个初始值 b_i ，希望最终将这个值从 b_i 变成 c_i

可以执行无限次操作，每次操作需要有两个值 u k

表示从以节点 u 为根的子树中选出 k 个节点随意交换它们的 b_i

其中 k 至少为 1 ，最多为以 u 为根的子树的节点个数

此时的花费为 k×a_u

如果存在操作方案，求出最小花费，不存在输出 -1

限制

Time limit per test: 2 seconds

Memory limit per test: 256 megabytes

• 1≤n≤2?10⁵
• 1≤ai≤10⁹ , 0≤bi,ci≤1
• 1≤u,v≤n, u≠v

解题思路

对于不存在的情况

如果 b 数组中 1 的个数与 c 数组中的个数对不上，则无法通过交换将每个节点从状态 b 变成状态 c

---

容易想到，我们应该从花费最小的节点开始，以最优方案去交换这个节点的子树

所以首先根据花费来将节点进行排序，然后开始处理子树

对于子树的搜索，用双向存图方案的话需要先求出每个节点的深度dep，然后根据深度大小判断是否为子树的节点；用单向存图方案不需要

---

处理子树时，需要遍历子树的所有节点——

如果某个节点 i 的 b_i 与 c_i 相同，则无需进行交换

如果不同，则肯定是 1->0 或者 0->1 的情况，使用 cntb 和 cntc 来储存这两种情况

处理完节点 i 的子树后，

cntb 则表示子树中尚未交换但需要交换的节点需要 1->0 的数量

cntc 则表示子树中尚未交换但需要交换的节点需要 0->1 的数量

可以发现，如果选择了两个节点，一个是 0->1 ，一个是 1->0

那么让他们互相交换原来的值就可以满足题意

此时对应的是从 cntb 中取出 1，从 cntc 中取出 1，凑成一对

所以总共可以执行交换的节点对数为 min(cntb,cntc) ，数量为 2*min(cntb,cntc)

花费则为 a_i * 2 * min(cntb,cntc)

但需要注意，可能此时会出现多出来一些 0->1 或者 1->0 的节点

所以需要再引入一个pair数组 les ，表示某个节点被访问过后还剩余的未处理的节点数，之后如果再次搜索到该节点时直接取 les 数组的值即可，同时也不用再返回去修改子树的值，做到记忆化搜索

---

因为有了les数组，所以如果按照花费从小到大排序后遍历

遍历过程中，如果某个节点已经被访问过，则不需要处理

如果没有访问过，搜索这个节点的子树，但搜索过程中如果遇到了访问过的节点，不能直接跳过

此时需要注意将访问过的节点的 les 数组值 加到 cntb 和 cntc 中

完整程序

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int,int> P;

struct node
{
	ll a;
	int id;
	bool operator < (const node& x) const
	{
		return a<x.a; //根据花费从小到大排序
	}
}arr[200050];
int b[200050],c[200050];
vector<int> G[200050]; //存图

bool vis2[200050];
int dep[200050]; //节点深度
void dfs2(int st,int d) //dfs2求每个节点的深度
{
	dep[st]=d;
	vis2[st]=true;
	for(int it:G[st])
	{
		if(!vis2[it])
			dfs2(it,d+1);
	}
}
bool vis[200050];
P les[200050];
/*
（只有当i已经被访问过时，即vis[i]==true时）
first表示处理完节点i的子树后，在b数列中剩余的1的数量
second表示处理完节点i的子树后，在c数列中剩余的1的数量
*/

int cntb,cntc;
void dfs(int p)
{
	if(vis[p]) //如果节点p已经被访问过
	{
		cntb+=les[p].first; //直接取尚未处理的种类数量加到cntb和cntc内即可
		cntc+=les[p].second;
		les[p].first=les[p].second=0; //置零
		return;
	}
	vis[p]=true;
	
	for(int it:G[p]) //搜索节点p的子树
	{
		if(!vis[it]) //如果没有访问过，就继续深搜
			dfs(it);
		else //如果已经被访问过了
		{
			cntb+=les[it].first; //直接取尚未处理的种类数量加到cntb和cntc内即可
			cntc+=les[it].second;
			les[it].first=les[it].second=0; //置零
		}
	}
	
	if(b[p]!=c[p]) //判断节点p自己
	{
		if(b[p]==1)
			cntb++;
		else
			cntc++;
	}
}

void solve()
{
	int n,tmp[2]={0,0},u,v;
	
	cin>>n;
	for(int i=1;i<=n;i++)
	{
		cin>>arr[i].a>>b[i]>>c[i];
		arr[i].id=i;
		
		if(b[i]==1)
			tmp[0]++;
		if(c[i]==1)
			tmp[1]++;
	}
	for(int i=1;i<n;i++)
	{
		cin>>u>>v;
		G[u].emplace_back(v);
		G[v].emplace_back(u);
	}
	
	if(tmp[0]!=tmp[1]) //如果b中1的个数与c中1的个数不同，说明无法通过交换得到c的状态
	{
		cout<<"-1
";
		return;
	}
	
	sort(arr+1,arr+n+1); //按照花费来排序
	dfs2(1,1); //求出每个节点的深度，1为根
	
	ll ans=0;
	for(int i=1;i<=n;i++) //从最小花费开始
	{
		int cur=arr[i].id;
		if(!vis[cur]) //只要这个节点没被访问过，才需要计算
		{
			vis[cur]=true;
			les[cur]=P(0,0);
			
			cntb=cntc=0;
			for(int it:G[cur]) //从与cur相邻的节点中寻找
				if(dep[it]>dep[cur]) //如果深度比cur大，说明是cur的子树
					dfs(it); //从子树中深搜
			
			if(b[cur]!=c[cur]) //判断cur自己
			{
				if(b[cur]==1)
					cntb++;
				else
					cntc++;
			}
			
			les[cur].first+=cntb;
			les[cur].second+=cntc;
			
			int d=min(les[cur].first,les[cur].second); //取较小的值，然后就能交换2*d个节点
			ans+=arr[i].a*d*2; //需要的花费
			
			les[cur].first-=d; //剩余未被处理的数量
			les[cur].second-=d;
		}
	}
	cout<<ans<<‘
‘;
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
	solve();
    return 0;
}