AtCoder Beginner Contest 163 F.path pass i

AtCoder Beginner Contest 163 F.path pass i

题目链接：163 F.path pass i

错解： 比赛时想到了用子节点数去分类统计，后面把自己绕晕了，当时写的时候由于 dfs 时只能由上往下遍历，当时看样例时只过了单侧链树的样例，于是往上找最靠近当前结点且与当前结点颜色相同的结点深度深度作差后可用公式直接得出两节点之间可选方案数，但这样会重复计算且当前结点还能与其祖先结点组合，显然思路是错的。

解题思路： 正向讨论情况太复杂，且无法讨论当前结点与其祖先结点的情况组合，于是我们想到用容斥的思想。首先有一个知识点：一棵有 n 个结点的树上不重复的路径数为n*(n+1)/2，很好证明其实就是1+2+……+n。所以初始化时将树上所有的结点看成一种颜色，即ans[i] = n*(n+1)/2，然后有两种情况是需要 “斥” 的。如图所示的一棵树，小圈内的数为结点标号，边上的数为节点颜色。情况一属于黑圈部分，情况二属于红圈部分。

情况一(图中黑圈)： 不满足条件的部分夹在两个颜色相同的结点之间 (结点1和结点4)，且这两个颜色相同的结点中的其中一个是 所夹部分的父节点，则所夹部分也是树，计算出所夹部分的结点数用公式求出不满足条件部分的路径条数减去即可。

情况二(图中红圈)： 不满足条件的部分夹在两个颜色相同的结点 (结点2和结点7) 之间，但所夹部分是这两个结点的祖先，这样根据 dfs 序根本不可能统计到上面部分，但由图中可看出，统计出每个结点的子树结点数后我们可用总结点数 n 减去结点2和结点7的子树结点数留下的就是红圈部分结点数，公式求出路径减去即可。

代码实现： 另外一些解释和需要注意的点在代码中都注解了。

#include <bits/stdc++.h>
#define ll long long
#define INF 0x3f3f3f3f
#define mem( f, x ) memset( f, x, sizeof( f ) )
#define pii pair<int, int>
#define fi first
#define se second
#define mk(x, y) make_pair( x, y )
#define pk push_back
using namespace std;
const int M = 2*1e5 + 5;
const int N = 2*1e5 + 5;
int m, n;
int color[N];
int head[N], cnt;
//son:子节点数量 cnt_color[i]: 以颜色i为根的子树结点数量
int son[N], cnt_color[N]; 
ll ans[N];

struct eg{
    int to, pre;
    eg(){ to = pre = 0; }
    eg( int tt, int pp ){
        to = tt, pre = pp;
    }
}e[2*M];

void add( int x, int y ){
    e[++cnt] = eg( y, head[x] );
    head[x] = cnt;
}

void init( ){
    ll tmp = cal( n );
    for( int i = 1; i <= n; i++ ){
        head[i] = son[i] = cnt_color[i] = 0;
        ans[i] = tmp; //初始化总数
    }
    cnt = 0;
}

void dfs( int cur, int fa ){
    son[cur] = 1;
    ll cur_c = color[cur];
    ll pre = cnt_color[cur_c]; //pre：遍历到cur结点之前统计到了多少以颜色cur_c为根的结点
    for( int i = head[cur]; i; i = e[i].pre ){
        int to = e[i].to;
        if( to == fa ) continue;
        //注：这里的tmp不能直接用上面的pre替换，
        //一个子节点遍历完后，可能会改变sum[cur_c]的值
        ll tmp = cnt_color[cur_c];
        dfs( to, cur );
        son[cur] += son[to];
        ll dt = cnt_color[cur_c]-tmp;
        ans[cur_c] -= cal( son[to]-dt );  //减去情况一
    }
    //cur的子节点必定以颜色cur_c为根，前面统计到的pre加上新加的son[cur]更新sum[cur]
    cnt_color[cur_c] = pre+son[cur];
}

int main( ){
    while( scanf( "%d", &n ) != EOF ){
        init( );
        for( int i = 1; i <= n; i++ )
            cin >> color[i];
        int x, y;
        for( int i = 1; i < n; i++ ){
            cin >> x >> y;
            add( x, y );
            add( y, x );
        }
        dfs( 1, 0 );
        for( int i = 1; i <= n; i++ ){
            ans[i] -= cal( n-cnt_color[i] ); //减去情况二
            cout << ans[i] << endl;
        }
    }
    return 0;
}