Codeforces Round #646 (Div. 2) E. Tree Shuffling (树上dfs统计)

E. Tree Shuffling

题意

有一棵 n 个结点,根为结点 1 的树,每个结点有一个选取代价 a i a_i ai,当前 b i b_i bi(0或1),目标数字 c i c_i ci(0或1) 。每次可以选择以一个结点为根节点的子树中的 k 个结点,然后任意分配它们的 b i b_i bi,消耗的代价为 k × a r o o t k×a_{root} k×aroot ,判断能否把所有结点都变为目标数字以及最小代价是多少。

思路

如果 b 中 1 的个数和 c 中 1 的个数不相同,那么肯定不可以,否则一定可行。

对于任意一个子树而言,它根节点的代价应该是它本身的代价以及所有父节点代价中的最小代价。

换言之,我们把根节点的代价向其子树传递,因为对于一个子树中如果有k个要交换的节点,并且这个子树根的父节点(祖先节点)中有代价小于它的,那么我们完全可以选择一个更大的子树,它的根节点消耗代价更少,所以需要将代价向下传递,每个节点的代价都是它所有父节点中最少的。

然后我们记录每个子树中需要交换的1和0的个数,进行交换,累加答案,有多余的1或0,则向上传递。

代码

#include

using namespace std;

const int N = 2e5+10;

int n;
int a[N];
int b[N];
int c[N];
vector<int> g[N];
long long ans;
int cnt[N][2];
void dfs(int u, int pre)
{
    if(pre!=0)
        a[u] = min(a[u],a[pre]);//代价向下传递
    for(auto v : g[u]){
        if(v==pre)
            continue;
        dfs(v,u);
        //多余1或0向上传递
        cnt[u][0] += cnt[v][0];
        cnt[u][1] += cnt[v][1];
    }
    if(b[u]!=c[u])//记录需要交换的1和0数量
        ++cnt[u][b[u]];
    int mi = min(cnt[u][0],cnt[u][1]);
    ans += 2LL * mi * a[u];
    cnt[u][0] -= mi;
    cnt[u][1] -= mi;
}

int main(){
    cin >> n;
    int b1 = 0;
    int c1 = 0;
    for(int i = 1; i <= n; ++i){
        cin >> a[i] >> b[i] >> c[i];
        if(b[i])
            ++b1;
        if(c[i])
            ++c1;
    }
    for(int i = 1; i < n; ++i){
        int u,v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    if(b1!=c1){
        cout << -1 << endl;
    }
    else{
        dfs(1,0);
        cout << ans << endl;
    }
    return 0;
}

你可能感兴趣的:(Codeforces)