【AtCoder】abc_198E - Unique Color【DFS】

题目链接
题意：给你一棵树，每个节点有一个颜色值，问满足条件（从节点$1$到节点$x$的路径上没有与$x$颜色相同的节点）的节点$x$有几个。
思路：dfs搜索，设置一个“从1号节点到当前节点的路径上每种颜色的顶点数量”的数组/map存储，并在每次进入/退出顶点时对其进行更新。更新复杂度$O(1)$，每一个节点只会遍历一次，总复杂度$O(n)$。注意用map存颜色记录，一开始用set存T到无语。
AC代码：

#include <bits/stdc++.h>
#define pcc pair<char, char>
#define pii pair<int, int>
#define vi vector<int>
#define vl vector<ll>
#define rep(i, x, y) for (int i = x; i <= y; i++)
#define per(i, x, y) for (int i = x; i >= y; i--)
#define rep0(i, n) for (int i = 0; i < (n); i++)
#define per0(i, n) for (int i = (n)-1; i >= 0; i--)
#define mp make_pair
#define pb push_back
#define F first
#define S second
#define sz(x) (x).size()
#define all(x) (x).begin(), (x).end()
#define ll long long
#define ull unsigned long long
#define db double
#define ld long double
using namespace std;

inline int read() {
  int x = 0, f = 1;
  char ch = getchar();
  while (ch < '0' || ch > '9') {
    if (ch == '-') f = -1;
    ch = getchar();
  }
  while (ch >= '0' && ch <= '9') {
    x = x * 10 + ch - '0';
    ch = getchar();
  }
  return x * f;
}

const double eps = 1e-9;
const int inf = 0x3f3f3f3f;
const int mod = 1e9 + 7;
const int maxn = 1e5 + 10;
const int MAX = 1e5;
const double pi = acos(-1.0);
const ll INF = 0x3f3f3f3f3f3f3f;
ll qpow(ll a, ll b) {
  ll res = 1;
  while (b) {
    if (b & 1) res = (res * a) % mod;
    a = (a * a) % mod;
    b >>= 1;
  }
  return res % mod;
}
/***************main****************/
ll T = 1;
ll m, n;
ll color[maxn], vis[maxn];
bool ans[maxn];
vector<ll> G[maxn];
unordered_map<int, int> ma;
void dfs(int x) {
  if (ma[color[x]] == 1) ans[x] = 1;
  for (auto i : G[x]) {
    if (!vis[i]) {
      vis[i] = 1;
      ma[color[i]]++;
      dfs(i);
      ma[color[i]]--;
    }
  }
}

int main() {
  n = read();
  rep(i, 1, n) color[i] = read();
  rep(i, 1, n - 1) {
    int x = read(), y = read();
    G[x].push_back(y);
    G[y].push_back(x);
  }
  vis[1] = 1;
  ans[1] = 1;
  ma[color[1]]++;
  dfs(1);
  rep(i, 1, n) if (ans[i]) printf("%d\n", i);
  return 0;
}