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;
}