hdu6035

树形dp

并没有想出来。。。

这个点分治很明显做不了，没办法批量处理信息，那么我们就只能考虑树形dp了，然后觉得用虚树什么的每个颜色单独考虑，GG

事实上我们可以这样考虑，我们对树进行dfs，每个节点有一个颜色，那么我们考虑子树中有多少条路径不包含这个颜色，很明显这条路径不能路过这个点，也就是说答案是C(n,2),n是每个儿子中不包含这个颜色的连通块大小，具体看程序，这里我们利用线段树合并来计算连通块大小，也可以用差分(?)统计一下就行了，在1号节点时要把除了1号节点的颜色也算上，然后用总路径条数减去即可

大佬还是强啊。。。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 200010;
vector<int> G[N];
int n, cnt, color;
ll ans;
int size[N], sum[N * 30], lc[N * 30], rc[N * 30], root[N], c[N], vis[N];
namespace Segment_Tree
{
    void update(int l, int r, int &o, int pos, int d)
    {
        o = ++cnt;
        sum[o] += d;
        if(l == r) return;
        int mid = (l + r) >> 1;
        if(pos <= mid) update(l, mid, lc[o], pos, d);
        else update(mid + 1, r, rc[o], pos, d);
    }
    int merge(int u, int v)
    {
        if(!u) return v;
        if(!v) return u;
        lc[u] = merge(lc[u], lc[v]);
        rc[u] = merge(rc[u], rc[v]);
        sum[u] += sum[v];
        return u;
    }
    int query(int l, int r, int o, int pos)
    {
        if(!o) return 0;
        if(l == r) return sum[o];
        int mid = (l + r) >> 1;
        if(pos <= mid) return query(l, mid, lc[o], pos);
        else return query(mid + 1, r, rc[o], pos);
    }
} using namespace Segment_Tree;
void dfs(int u, int last)
{
    size[u] = 1;
    ll tot = 0;    
    for(int i = 0; i < G[u].size(); ++i)
    {
        int v = G[u][i];
        if(v == last) continue;
        dfs(v, u); 
        size[u] += size[v];
        ll tmp = query(1, n, root[v], c[u]);
        tot += tmp;
        ans -= (size[v] - tmp) * (size[v] - tmp - 1) / 2ll;
    }
    update(1, n, root[u], c[u], size[u] - tot);
    for(int i = 0; i < G[u].size(); ++i) 
    {
        int v = G[u][i];
        if(v == last) continue;
        root[u] = merge(root[u], root[v]);
    }
} 
int main()
{
    scanf("%d", &n);
    ans = (ll)n * (ll)(n - 1) / 2ll; 
    for(int i = 1; i <= n; ++i) scanf("%d", &c[i]), vis[c[i]] = 1;
    for(int i = 1; i <= n; ++i) color += vis[i];
    ans *= color;
    for(int i = 1; i < n; ++i)
    {
        int u, v;
        scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1, 0);
    for(int i = 1; i <= n; ++i) if(vis[i] && c[1] != i)
    {
        ll tmp = query(1, n, root[1], i);
        ans -= (ll)(size[1] - tmp) * (ll)(size[1] - tmp - 1) / 2ll;
    }
    printf("%lld
", ans);
    return 0;
}