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