HDU 6035 (虚树)(统计颜色)

HDU 6035 Colorful Tree

Problem : 给一棵树，每一个结点有一种颜色，定义每条路径的权值为这条路径上颜色的种数，询问全部路径(C(n,2)条)的权值之和。
Solution : 分开考虑每种颜色对答案的贡献。对于一种颜色，一共有C(n,2)条路径，再考虑问题的反面。不含有这种颜色的路径的数量，即为将这棵树上全部这种颜色的点删去后，全部剩下的树的路径数量。
问题关键在于统计删去点以后每颗树的大小，对于每一个节点要计算删去其全部相同颜色的子树，用栈进行维护。即每次访问到每一个节点时，将其最近的相同颜色的祖先的大小减去该子树的大小。

#include <bits/stdc++.h>

using namespace std;

const int N = 200008;

vector <int> eg[N];

int cl[N];
int flag[N];
int size[N];
int nt[N], st[N];
long long tag[N];
long long tagrt[N];

int n;

void dfs(int u, int fa) 
{
    size[u] = 1;
    for (auto v : eg[u])
    {
        if (v == fa) continue;
        dfs(v, u);
        size[u] += size[v];
    }
}
void solve(int u, int fa)
{
    int rt;
    int last = st[cl[u]];
    if (last == 0) 
    {
        rt = 1; 
        tagrt[cl[u]] += size[u];
    }
    else 
    {
        rt = nt[last];
        tag[rt] += size[u];
    }
    st[cl[u]] = u;
    for (auto v: eg[u])
    {
        if (v == fa) continue;
        nt[u] = v;
        solve(v, u);
    }
    st[cl[u]] = last;
}

int main()
{
    cin.sync_with_stdio(0);
    int cas = 0;
    while (cin >> n)
    {
        for (int i = 1; i <= n; ++i) 
        {
            size[i] = nt[i] = tag[i] = tagrt[i] = flag[i] = 0, eg[i].clear();
        }
        for (int i = 1; i <= n; ++i) cin >> cl[i], flag[cl[i]] = 1;
        int num = 0;
        for (int i = 1; i <= n; ++i) num += flag[i];
        long long ans = 1ll * n * (n - 1) / 2 * num;
        for (int i = 1; i < n; ++i)
        {
            int u, v;
            cin >> u >> v;
            eg[u].push_back(v);
            eg[v].push_back(u);
        }
        dfs(1, 0);
        solve(1, 0);

        for (int i = 2; i <= n; ++i)
        {
            int p = size[i] - tag[i];
            ans -= 1ll * p * (p - 1) / 2;
        }
        for (int i = 1; i <= n; ++i)
            if (flag[i]) 
            {
                int p = size[1] - tagrt[i];
                ans -= 1ll * p * (p - 1) / 2;
            }
        cout << "Case #" << ++cas << ": " << ans << endl;
    }
}