Educational Codeforces Round 121 (Rated for Div. 2) E. Black and White Tree

题目链接

题意

一棵n个节点的树,节点有黑白两种颜色,黑点的数量至少为2。
一开始有一个石子在树的某一个点上,每一次操作你可以选择
一个黑点,让石子向黑点的方向走一条边。不能连续两次操作
选择相同的黑点。对于每一个点,求出如果石子在这个点上,是
否有方法可以将石子移动到黑点上。

思路：

如何确定当前i是否满足题意？
（1）对于i来说，本身就是黑色，则i满足
（2）对于i不是黑色来说，他的子树中存在2个及以上的黑色结点，则i满足，同时i这个结点对于他的父节点来说，可以连续访问，并成立。
如何实现呢？
（1）计算每个结点的子树有多少黑色的结点
（2）判断从i结点开始，是否存在一个黑色结点j，使得i->j。

代码

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN=3e5+5;
vector<int>edges[MAXN];
int c[MAXN];
int sz[MAXN];//sz[i]表示以i为根，子树中黑色的颜色个数为sz[i]
int vis[MAXN];//vis[i]表示从i开始是否满足题意
void dfs1(int u,int fa){
    sz[u]=vis[u]=c[u];//如果u为黑色，则满足题意
    for(auto& v:edges[u]){
        if(v!=fa){
            dfs1(v,u);
            sz[u]+=sz[v];
            if(c[v]||(vis[v]&&sz[v]>=2))vis[u]=1;//如果当前结点为黑色，则父节点满足题意，如果当前结点满足题意，并且子树有两个或两个以上的黑色结点，则父节点满足题意
        }
    }
}
void dfs2(int u,int fa){
    for(auto& v:edges[u]){
        if(v==fa)continue;
        if(c[u]||(vis[u]&&sz[1]-sz[v]>=2))vis[v]=1;//如果父节点为黑色，则孩子结点v可以满足题意，如果父节点满足题意，并且不包括当前子树的黑色结点个数，还有2个及以上，则满足题意
        dfs2(v,u);
    }
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    int n;cin>>n;
    for(int i=1;i<=n;i++)cin>>c[i];
    for(int i=1;i<n;i++){
        int u,v;cin>>u>>v;
        edges[u].push_back(v);
        edges[v].push_back(u);
    }
    dfs1(1,0);
    dfs2(1,0);
    for(int i=1;i<=n;i++)cout<<vis[i]<<" ";cout<<endl;
    return 0;
}