mobius HDOJ 5468 Puzzled Elena

dfs遍历树+莫比乌斯反演..

#include <bits/stdc++.h>using namespace std;typedef long long LL;const int maxn = 100005;const int maxm = 200005;struct Edge{int v;Edge *next;}E[maxm], *H[maxn], *edges;int not_prime[maxn];int prime[maxn];int mu[maxn];int p_cnt;vector<int> vec[maxn], v[maxn];int cnt[maxn];int res[maxn];int a[maxn];int n;void Init(){mu[1] = 1;p_cnt = 0;for(int i = 2; i < maxn; i++) {if(!not_prime[i]) {prime[p_cnt++] = i;mu[i] = -1;}for(int j = 0; j < p_cnt && i * prime[j] < maxn; j++) {not_prime[i * prime[j]] = 1;if(i % prime[j]) mu[i * prime[j]] = -mu[i];else {mu[i * prime[j]] = 0;break;}}}}void init(){edges = E;memset(H, 0, sizeof H);memset(cnt, 0, sizeof cnt);memset(res, 0, sizeof res);}void addedges(int u, int v){edges->v = v;edges->next = H[u];H[u] = edges++;}int calc(int u){int ans = 0;for(int i = 0; i < vec[u].size(); i++) {int d = vec[u][i];ans += cnt[d] * mu[d];}return ans;}void dfs(int u, int fa){res[u] -= calc(u);for(int i = 0; i < vec[u].size(); i++) {int d = vec[u][i];cnt[d]++;}for(Edge *e = H[u]; e; e = e->next) if(e->v != fa) dfs(e->v, u);res[u] += calc(u);}void work(){for(int i = 1; i < n; i++) {int u, v;scanf("%d%d", &u, &v);addedges(u, v);addedges(v, u);}for(int i = 1; i <= n; i++) scanf("%d", &a[i]);for(int i = 1; i <= n; i++) vec[i] = v[a[i]];dfs(1, 1);for(int i = 1; i <= n; i++) printf("%d%c", res[i], i == n ? '\n' : ' ');}int main(){Init();int _ = 0;for(int i = 1; i < maxn; i++)for(int j = i; j < maxn; j += i)v[j].push_back(i);while(scanf("%d", &n) != EOF) {printf("Case #%d: ", ++_);init();work();}return 0;}