hdu 6035-树状DP

题目链接：点击打开链接

题解思路：一开始绞尽脑汁怎么正向解决这个问题，就是定义一个优先级然后看每个节点在优先级的情况下的贡献点是多少，发现好难做啊。看了别人的博客之后发现反向思考原来简单的多。假设这颗树的所以不用颜色有sums种，那么我们假设每种颜色对每个路径都有贡献那么答案应该就是（所以路径数*颜色数）- （所以颜色没有参与贡献的总和）。那么我们用sum数组保存，sum【i】表示i颜色能参与到贡献的节点数，每次递归更新的时候保存现在状态的sum【i】 然后用更新后的减去现在的就是新增的参与贡献数，用cnt【x】- (新增的数) 就是区域内未参与贡献节点数.

代码：

#include<algorithm>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<vector>
typedef long long ll;
using namespace std;

const int mx = 2e5+10;
int n,m,c[mx],cnt[mx],size;
ll ans,sum[mx];
bool vis[mx];
vector<int> vec[mx];
int dfs(int x,int fa){
    cnt[x] = 1;//自己包括子树的节点数和
    ll num = 0;
    for(int i=0;i<vec[x].size();i++){
        int son = vec[x][i];
        if(son == fa) continue;
        ll last = sum[c[x]]; // 当前状态下此节点的子树被C[x]覆盖了多少
        cnt[x] += dfs(son,x);
        ll add = sum[c[x]] - last; // 递归前后的变化就是没有参与贡献的节点数
        ans += (cnt[son]-add)*(cnt[son]-add-1)/2;
        num += cnt[son] - add; // 子节点部分新增加的覆盖数量
    }
    sum[c[x]] += num + 1;
    return cnt[x];
}
int main(){
    int t,cas=1,sums;
    while(~scanf("%d",&n)){
        memset(vis,0,sizeof(vis));
        sums = ans = size = 0;
        for(int i=1;i<=n;i++){
             scanf("%d",c+i);
             if(!vis[c[i]]) sums++;
             vis[c[i]] = 1;
             vec[i].clear();
             sum[i] = 0;
        }
        int x,y;
        for(int i=1;i<n;i++){
            scanf("%d%d",&x,&y);
            vec[x].push_back(y);
            vec[y].push_back(x);
        }
        dfs(1,0);
        for(int i=1;i<=n;i++){
            if(!vis[i]) continue;
            ans += (n-sum[i])*(n-sum[i]-1)/2; // 剩余没有贡献部分（不被夹在内部的部分）
        }
        printf("Case #%d: %lld\n",cas++,1ll*n*(n-1)/2*sums-ans);
    }
    return 0;
}