AtCoder Beginner Contest 359 G Solution

题面

here

思路

感觉十分眼熟。

求同色点的树上距离和。

拆成每条边对答案的贡献。

把一条边切断，将树分成两部分。

同理，同色的点也被分成两部分。

那么一条边对答案的贡献就是所有同色点两边数量乘积的和。

解法

启发式合并。

使用 map 存子树内每一种颜色的个数，边合边算。

时间复杂度 \(O(n\log n)\)。

代码

#include <iostream>
#include <map>
#include <vector>
using namespace std;
const int N = 2e5 + 10;
using mii = map<int, int>;
int n, a[N];
mii mp[N];
using ll = long long;
ll cnt[N], res[N], ret;
vector<int> road[N];
void merge(int x, int y)
{
    if (mp[x].size() < mp[y].size())
    {
        mp[x].swap(mp[y]);
        swap(res[x], res[y]);
    }
    for (auto &[k, v] : mp[y])
    {
        res[x] -= mp[x][k] * (cnt[k] - mp[x][k]);
        mp[x][k] += v;
        res[x] += mp[x][k] * (cnt[k] - mp[x][k]);
    }
}
void dfs(int x, int f)
{
    mp[x][a[x]] = 1;
    res[x] = cnt[a[x]] - 1;
    for (auto &i : road[x])
    {
        if (i == f)
            continue;
        dfs(i, x);
        merge(x, i);
    }
    ret += res[x];
}
int main()
{
    scanf("%d", &n);
    for (int i = 1, x, y; i < n; i++)
    {
        scanf("%d%d", &x, &y);
        road[x].push_back(y);
        road[y].push_back(x);
    }
    for (int i = 1; i <= n; i++)
    {
        scanf("%d", a + i);
        cnt[a[i]]++;
    }
    dfs(1, 0);
    printf("%lld\n", ret);
}