Skip to content

Instantly share code, notes, and snippets.

@chermehdi
Created July 18, 2019 18:39
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save chermehdi/e0a48ead9aa208f2ac96132deb8ddf67 to your computer and use it in GitHub Desktop.
Save chermehdi/e0a48ead9aa208f2ac96132deb8ddf67 to your computer and use it in GitHub Desktop.
#include <algorithm>
#include <iostream>
#include <map>
#include <set>
#include <unordered_map>
#include <vector>
using namespace std;
const int N = 100005;
int cnt[N], ans[N], val[N], size[N], tree[N];
bool is_big[N];
vector<int> g[N];
int n, m;
void update(int index, int val) {
while (index < N) {
tree[index] += val;
index += (index & (-index));
}
}
int query(int index) {
int ret = 0;
while (index > 0) {
ret += tree[index];
index -= (index & -index);
}
return ret;
}
void compress() {
unordered_map<int, int> mp;
vector<int> cp(n);
int pt = 1;
for (int i = 1; i <= n; ++i) {
if (mp[val[i]]) {
cp[i] = mp[val[i]];
} else {
mp[val[i]] = pt;
cp[i] = pt++;
}
}
for (int i = 1; i <= n; ++i) {
val[i] = cp[i];
}
}
void add(int u, int p, int x) {
if (cnt[val[u]] == 0) {
update(val[u], 1);
}
cnt[val[u]] += x;
if (cnt[val[u]] == 0) {
update(val[u], -1);
}
for (int v : g[u]) {
if (v != p && !is_big[v]) {
add(v, u, x);
}
}
}
void dfs(int u, int p) {
size[u] = 1;
for (int v : g[u]) {
if (v != p) dfs(v, u);
size[u] += size[v];
}
}
void dsu(int u, int p, bool keep) {
int big = 0;
for (int v : g[u]) {
if (v != p && size[big] < size[v]) {
big = v;
}
}
for (int v : g[u]) {
if (v != p && v != big) {
dsu(v, u, false);
}
}
if (big != 0) {
dsu(big, u, true);
is_big[big] = true;
}
add(u, p, 1);
ans[u] = query(N - 1);
is_big[big] = false;
if (!keep) {
add(u, p, -1);
}
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
int root;
cin >> n >> m >> root;
for (int i = 0; i < n - 1; ++i) {
int u, v;
cin >> u >> v;
g[u].push_back(v);
g[v].push_back(u);
}
for (int i = 1; i <= n; ++i) {
cin >> val[i];
}
dfs(root, 0);
compress();
dsu(root, 0, true);
for (int i = 0; i < m; ++i) {
int u;
cin >> u;
cout << ans[u] << '\n';
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment