HDU 6035 (虚树)(统计颜色)-优快云博客

本文解析了HDU6035 ColorfulTree问题，介绍了一种通过计算每种颜色对答案贡献的方法来求解所有路径权值之和。使用DFS遍历树结构，并利用栈维护每种颜色节点的祖先信息，最终得出不含特定颜色的所有子树的大小。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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;
    }
}

转载于:https://www.cnblogs.com/rpSebastian/p/7236153.html