HDU 4812 D Tree

#include<queue>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long LL;
const int INF = 0x7FFFFFFF;
const int mod = 1e6 + 3;
const int maxn = 2e5 + 10;
int n, K, x, y, inv[mod], h[mod], f[mod];

struct Tree
{
	int ft[maxn], nt[maxn], u[maxn], v[maxn], sz, n;
	int vis[maxn], cnt[maxn], mx[maxn], a, b, flag;
	void clear(int n)
	{
		this->n = n;	mx[0] = INF;
		a = b = sz = flag = 0;
		memset(h, 0, sizeof(h));
		for (int i = 1; i <= n; i++)
		{
			ft[i] = -1;
			vis[i] = 0;
			scanf("%d", &v[i]);
		}
	}
	void AddEdge(int x, int y)
	{
		u[sz] = y; nt[sz] = ft[x]; ft[x] = sz++;
	}
	int dfs(int x, int fa, int sum)
	{
		int ans = mx[x] = 0;
		cnt[x] = 1;
		for (int i = ft[x]; i != -1; i = nt[i])
		{
			if (vis[u[i]] || u[i] == fa) continue;
			int y = dfs(u[i], x, sum);
			if (mx[y]<mx[ans]) ans = y;
			cnt[x] += cnt[u[i]];
			mx[x] = max(mx[x], cnt[u[i]]);
		}
		mx[x] = max(mx[x], sum - cnt[x]);
		return mx[x] < mx[ans] ? x : ans;
	}
	void get(int x, int fa, int now, int kind)
	{
		if (!kind)
		{
			int y = (LL)K*inv[now] % mod;
			if (h[y] == flag)
			{
				if (a + b == 0) a = min(f[y], x), b = max(f[y], x);
				else
				{
					if (a > min(f[y], x)) a = min(f[y], x), b = max(f[y], x);
					else if (a == min(f[y], x) && b > max(f[y], x)) b = max(f[y], x);
				}
			}
		}
		else
		{
			if (h[now] != flag) h[now] = flag, f[now] = x;
			else f[now] = min(f[now], x);
		}
		for (int i = ft[x]; i != -1; i = nt[i])
		{
			if (vis[u[i]] || u[i] == fa) continue;
			get(u[i], x, (LL)now*v[u[i]] % mod, kind);
		}
	}
	void find(int x)
	{
		h[1] = ++flag;	f[1] = x;
		for (int i = ft[x]; i != -1; i = nt[i])
		{
			if (vis[u[i]]) continue;
			get(u[i], x, (LL)v[u[i]] * v[x] % mod, 0);
			get(u[i], x, v[u[i]], 1);
		}
	}
	void work(int x, int sum)
	{
		int y = dfs(x, -1, sum);
		vis[y] = 1;	 find(y);	
		for (int i = ft[y]; i != -1; i = nt[i])
		{
			if (vis[u[i]]) continue;
			if (cnt[u[i]] > cnt[y]) cnt[u[i]] = sum - cnt[y];
			work(u[i], cnt[u[i]]);
		}
	}
}solve;

void read(int &x)
{
	char ch;
	while ((ch = getchar()) < '0' || ch > '9');
	x = ch - '0'; 
	while ((ch = getchar()) >= '0' && ch <= '9') x = x * 10 + ch - '0';
}

void init()
{
	inv[0] = 0;	inv[1] = 1;
	for (int i = 2; i < mod; i++) inv[i] = (LL)inv[mod%i] * (mod - mod / i) % mod;
}

int main()
{
	init();
	while (~scanf("%d%d", &n, &K))
	{
		solve.clear(n);
		for (int i = 1; i < n; i++)
		{
			scanf("%d%d", &x, &y);
			solve.AddEdge(x, y);
			solve.AddEdge(y, x);
		}
		solve.work(1, n);
		if (solve.a + solve.b) printf("%d %d\n", solve.a, solve.b);
		else printf("No solution\n");
	}
	return 0;
}