Codeforces Round #635C Linova and Kingdom 思维

Linova and Kingdom

题意

现在有一颗n个节点的树，每个节点是一个城市，现在要选出k个城市作为工业城市，其他城市作为旅游城市，现在每个工业城市要派出一名特使前往根节点，每个特使的幸福度为经过的旅游城市的数量，求最大的幸福度总和。

思路

对于某个节点u，如果u是工业城市，那么它的子节点肯定是工业城市。

如果它的某个子节点不是，我们完全可以把它的子节点作为工业城市，而不是u。

我们再看如果选择了u作为工业城市，幸福度发生的变化。

假如本来总幸福度是ans

定义根节点的深度为1，u是不是工业城市只会对它的子树有影响

ans=ans-(sum[u]-1)*dep[u]先减去u为旅游城市时，子树的幸福度

ans=ans+sum[u]*(dep[u]-1)加上u变为工业城市时，子树的幸福度

化简一下:

ans=ans+dep[u]-sum[u]

如果u节点作为工业城市，它的贡献为dep[u]-sum[u]，即深度-子树大小

我们按照每个节点的深度-子树大小从大到小排序，取前k个。

代码

#include<bits/stdc++.h>
#define pb push_back
using namespace std;
typedef long long ll;
const int N=1e6+10;

vector<int>xiao[N],vec;
int dep[N],sum[N];
void dfs(int u,int pre)
{
    sum[u]=1;
    for(int v:xiao[u])
    {
        if(v==pre) continue;
        dep[v]=dep[u]+1;
        dfs(v,u);
        sum[u]+=sum[v];
    }
}
bool cmp(int a,int b)
{
    return (dep[a]-sum[a])>(dep[b]-sum[b]);
}
int main()
{
    int n,k;
    scanf("%d%d",&n,&k);
    for(int i=1;i<n;i++)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        xiao[u].pb(v);
        xiao[v].pb(u);
    }
    for(int i=1;i<=n;i++)
        vec.pb(i);
    dep[1]=1;
    dfs(1,1);
    sort(vec.begin(),vec.end(),cmp);
    ll maxn=0;
    for(int i=0;i<k;i++)
    {
        int u=vec[i];
        maxn+=dep[u]-sum[u];
    }
    printf("%lld\n",maxn);
    return 0;
}