题目链接
题意:给你一棵树,每个节点有一个颜色值,问满足条件(从节点
1
1
1到节点
x
x
x的路径上没有与
x
x
x颜色相同的节点)的节点
x
x
x有几个。
思路:dfs搜索,设置一个“从1号节点到当前节点的路径上每种颜色的顶点数量”的数组/map存储,并在每次进入/退出顶点时对其进行更新。更新复杂度
O
(
1
)
O(1)
O(1),每一个节点只会遍历一次,总复杂度
O
(
n
)
O(n)
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;
}