树 + 子树信息的维护

题目链接：http://codeforces.com/contest/1118/problem/F1
题目大意：有一棵树，有的节点是红色，有的节点是蓝色，有的节点没有颜色。把一条边拆开，分成两棵树，如果这两棵树都没有同时拥有两种颜色的节点那么这条边就是“漂亮的边”。问有多少这样的边。

思路：以dfs的顺序维护dp[i][0], dp[i][1]。分别表示以节点i为子树的红色节点数，和以节点i为子树的蓝色节点数。

dp[i][0]=节点i所有的子树的红色个数+(a[i]==1?1:0)/*自身的颜色*/

节点i所有的子树:dfs的顺序就是自己子树的根节点。
所以：dp[i][0]=∑(dp[i][j])+(a[i]==1?1:0)节点i到节点j有边相连。
并且节点j不是来到节点i的那个节点。那是节点i的父节点。

void dfs(int k)
{
    vis[k]=1;
    if(f[k]==1)        //初始化自身节点颜色
    {
        dp[k][0]++;
    }
    if(f[k]==2)
    {
        dp[k][1]++;
    }

    for(int i=e[k].size()-1;i>=0;i--)
    {
        int to=e[k][i];
        if(!vis[to])
        {
            dfs(to);
            dp[k][0]+=dp[to][0];//所有子节点的颜色
            dp[k][1]+=dp[to][1];
        }
    }
}

此时对于一条边，以dfs的顺序判断它的任意子树上的红色与蓝色的节点数。与另外一棵树的红色与蓝色的节点数。是否满足条件。
对于子树上的红色与蓝色的节点数就是dp[i][0], dp[i][1]。而另外一棵树可以用总的红色节点数-dp[i][0],总的 蓝色的节点数-dp[i][1]计算。

int  ans=0;
void DFS(int k)
{
    vis[k]=1;
    if(dp[k][0]==A&&dp[k][1]==0)//如果子树拥有全部红色节点
    {
        ans++;
    }
    else if(dp[k][1]==B&&dp[k][0]==0)//如果子树拥有全部蓝色节点
    {
        ans++;
    }
    for(int i=e[k].size()-1;i>=0;i--)
    {
        int to=e[k][i];
        if(!vis[to])
        {
            DFS(to);
        }
    }
}

思考：经典的维护子树信息的题。

全部代码：

#include<bits/stdc++.h>
#define LL long long
//#define p1 first
//#define p2 second
using namespace std;
const int mod=1e9+7;
//memset(a, 0, sizeof(a));
//next(p, 1), prev(p, 1), set<int>::iterator
//stack堆栈 queue队列 priority_queue优先队列
//vector向量 multiset平衡二叉树 deque双端队列
//pair{T1 first;T2 second;} greater<T>
//unordered_map 哈希map

vector<int> e[300005]; //存图
int f[300005];         //保存节点的颜色
int A, B;              //整棵树的红色与蓝色节点总和
int dp[300005][2]={0}; //以i节点为根节点的红色与蓝色节点数

int vis[300005]={0};
void dfs(int k)
{
    vis[k]=1;
    if(f[k]==1)        //初始化自身节点颜色
    {
        dp[k][0]++;
    }
    if(f[k]==2)
    {
        dp[k][1]++;
    }

    for(int i=e[k].size()-1;i>=0;i--)
    {
        int to=e[k][i];
        if(!vis[to])
        {
            dfs(to);
            dp[k][0]+=dp[to][0];//所有节点的颜色
            dp[k][1]+=dp[to][1];
        }
    }
}

int  ans=0;
void DFS(int k)
{
    vis[k]=1;
    if(dp[k][0]==A&&dp[k][1]==0)//如果子树拥有全部红色节点
    {
        ans++;
    }
    else if(dp[k][1]==B&&dp[k][0]==0)//如果子树拥有全部蓝色节点
    {
        ans++;
    }
    for(int i=e[k].size()-1;i>=0;i--)
    {
        int to=e[k][i];
        if(!vis[to])
        {
            DFS(to);
        }
    }
}

int main()
{
    int n;
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&f[i]);
    }
    for(int i=0;i<n-1;i++)
    {
        int x, y;
        scanf("%d%d",&x,&y);
        e[x].push_back(y);
        e[y].push_back(x);
    }
    A=count(f+1, f+1+n, 1);
    B=count(f+1, f+1+n, 2);
    dfs(1);
    memset(vis, 0, sizeof(vis));
    DFS(1);
    cout<<ans<<endl;

    return 0;
}

#include<bits/stdc++.h>
#define LL long long
//#define p1 first
//#define p2 second
using namespace std;
const int mod=1e9+7;
//memset(a, 0, sizeof(a));
//next(p, 1), prev(p, 1), set<int>::iterator
//stack堆栈 queue队列 priority_queue优先队列
//vector向量 multiset平衡二叉树 deque双端队列
//pair{T1 first;T2 second;} greater<T>
//unordered_map 哈希map

vector<int> e[300005]; //存图
int f[300005];         //保存节点的颜色
int A, B;              //整棵树的红色与蓝色节点总和
int dp[300005][2]={0}; //以i节点为根节点的红色与蓝色节点数

int vis[300005]={0};
void dfs(int k)
{
    vis[k]=1;
    if(f[k]==1)        //初始化自身节点颜色
    {
        dp[k][0]++;
    }
    if(f[k]==2)
    {
        dp[k][1]++;
    }

    for(int i=e[k].size()-1;i>=0;i--)
    {
        int to=e[k][i];
        if(!vis[to])
        {
            dfs(to);
            dp[k][0]+=dp[to][0];//所有节点的颜色
            dp[k][1]+=dp[to][1];
        }
    }
}

int DFS(int n)
{
    int ans=0;
    for(int i=1;i<=n;i++)
    {
        if(dp[i][0]==A&&dp[i][1]==0)     //如果子树拥有全部红色节点
        {
            ans++;
        }
        else if(dp[i][1]==B&&dp[i][0]==0)//如果子树拥有全部蓝色节点
        {
            ans++;
        }
    }

    return ans;
}

int main()
{
    int n;
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&f[i]);
    }
    for(int i=0;i<n-1;i++)
    {
        int x, y;
        scanf("%d%d",&x,&y);
        e[x].push_back(y);
        e[y].push_back(x);
    }
    A=count(f+1, f+1+n, 1);
    B=count(f+1, f+1+n, 2);
    dfs(1);
    memset(vis, 0, sizeof(vis));
    cout<<DFS(n)<<endl;

    return 0;
}