cf补题计划_Edu 162_E

E. Count Paths

题目简介

乍一眼一看是一个很简单树上dp（实际上也是树上dp
第一看做法是考虑dp[i][j]表达为第i个节点的下面有多少个颜色为j的节点，且这个节点于i节点之间无其他的节点为j颜色，然后dp转移十分的简单，但是n是$2\cdot {10}^5$,绝对超时

然后显然我们发现，一个节点只有一个颜色，根据时间复杂度来看的话，一个节点所储存的信息不会太多，所以我们考虑一个节点只存一个信息就是他自身颜色的信息，

不难发现对于节点$i$来说 他的颜色为$colo{r}_i$如果dfs进入该点，i节点的子树中的$colo{r}_i$点和子树外面的外界$colo{r}_i$点无法相连，而其他颜色并无影响，这样的话我们就可以用一个节点信息就表达出来了，所以考虑成立

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define pb push_back
#define pii pair<int, int>
#define all(a) a.begin(), a.end()
#define debug(x) cout << #x << " = " << x << endl
#define sz(a) ((int)a.size())
const int mod = 998244353, N = 2e5 + 10, G = 3;
vector<int> edge[N];
ll a[N], dp[N], sum[N], ans = 0;
void dfs(int u, int father)
{
    ans += sum[a[u]];      // u节点于color a[u]形成最美路径
    dp[u] = sum[a[u]] + 1; // u 节点子树外面的color a[u]的数目
    for (int i = 0; i < edge[u].size(); i++)
    {
        sum[a[u]] = 1; // 进入u节点的子树里面color a[u] 的数量变为一，与外界隔离
        int v = edge[u][i];
        if (v == father)
        {
            continue;
        }
        dfs(v, u);
    }
    sum[a[u]] = dp[u]; // 将sum a[u]变为子树外面的值
}
void solve()
{
    int n;
    cin >> n;
    ans = 0;
    for (int i = 1; i <= n; i++)
    {
        cin >> a[i];
        sum[i] = 0;
        edge[i].clear();
    }
    for (int i = 1; i < n; i++)
    {
        int u, v;
        cin >> u >> v;
        edge[u].pb(v);
        edge[v].pb(u);
    }
    dfs(1, 0);
    cout << ans << endl;
}

int main()
{
    ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
    int t = 1;
    cin >> t;
    while (t--)
    {
        solve();
    }
    return 0;
}