E. Count Paths
题目简介
乍一眼一看是一个很简单树上dp（实际上也是树上dp
第一看做法是考虑dp[i][j]表达为第i个节点的下面有多少个颜色为j的节点，且这个节点于i节点之间无其他的节点为j颜色，然后dp转移十分的简单，但是n是\(2\cdot10^5\),绝对超时
然后显然我们发现，一个节点只有一个颜色，根据时间复杂度来看的话，一个节点所储存的信息不会太多，所以我们考虑一个节点只存一个信息就是他自身颜色的信息，
不难发现对于节点\(i\)来说 他的颜色为\(color_i\)如果dfs进入该点，i节点的子树中的\(color_i\)点和子树外面的外界\(color_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]];dp[u] = sum[a[u]] + 1;for (int i = 0; i < edge[u].size(); i++){sum[a[u]] = 1;int v = edge[u][i];if (v == father){continue;}dfs(v, u);}sum[a[u]] = dp[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;
}
