【树分治】 HDOJ 4812 D Tree

比较浅显的树分治。。。对于根节点，找出两条路径使得乘积为k就行了。。。。但是这道题卡时间。。。要写的很好才行。。。

#include <iostream>
#include <queue>
#include <stack>
#include <map>
#include <set>
#include <bitset>
#include <cstdio>
#include <algorithm>
#include <cstring>
#include <climits>
#include <cstdlib>
#include <cmath>
#include <time.h>
#define maxn 100005
#define maxm 400005
#define eps 1e-12
#define mod 1000003
#define INF 0x3f3f3f3f
#define PI (acos(-1.0))
#define lowbit(x) (x&(-x))
#define mp make_pair
#define ls o<<1
#define rs o<<1 | 1
#define lson o<<1, L, mid 
#define rson o<<1 | 1, mid+1, R
#pragma comment(linker, "/STACK:102400000,102400000")
#define pii pair<int, int> 
typedef long long LL;
typedef unsigned long long ULL;
//typedef int LL;
using namespace std;
LL qpow(LL a, LL b){LL res=1,base=a;while(b){if(b%2)res=res*base;base=base*base;b/=2;}return res;}
LL powmod(LL a, LL b){LL res=1,base=a;while(b){if(b%2)res=res*base%mod;base=base*base%mod;b/=2;}return res;}
// head

struct Edge
{
	int v;
	Edge *next;
}E[maxm], *H[maxn], *edges;

pii h[mod + 5];
vector<pair<int, int> > dis;
int size[maxn];
bool done[maxn];
int mx[maxn];
int inv[mod + 5];
int a[maxn];
int kk, T;
pii ans;
int n, root, nsize;

void addedges(int u, int v)
{
	edges->v = v;
	edges->next = H[u];
	H[u] = edges++;
}

void init()
{
	T = 0;
	edges = E;
	memset(H, 0, sizeof H);
	memset(done, 0, sizeof done);
	memset(h, 0, sizeof h);
}

void getroot(int u, int fa)
{
	size[u] = 1, mx[u] = 0;
	for(Edge *e = H[u]; e; e = e->next) if(!done[e->v] && e->v != fa) {
		int v = e->v;
		getroot(v, u);
		size[u] += size[v];
		mx[u] = max(mx[u], size[v]);
	}
	mx[u] = max(mx[u], nsize - size[u]);
	if(mx[u] < mx[root]) root = u;
}

void getdis(int u, int fa, LL dist)
{
	dist = dist * a[u] % mod;
	dis.push_back(mp(dist, u));
	int t = (LL)kk * inv[dist] % mod;
	if(h[t].second == T) {
		ans = min(ans, mp(u, h[t].first));
		ans = min(ans, mp(h[t].first, u));
	}
	for(Edge *e = H[u]; e; e = e->next) if(!done[e->v] && e->v != fa) getdis(e->v, u, dist);
}

void solve(int u)
{
	T++;
	kk = (LL)kk * a[u] % mod;
	h[a[u]] = mp(u, T);
	done[u] = true;
	for(Edge *e = H[u]; e; e = e->next) if(!done[e->v]) {
		int v = e->v;
		dis.clear();
		getdis(v, u, a[u]);
		for(int i = 0; i < dis.size(); i++) {
			if(h[dis[i].first].second == T) h[dis[i].first] = mp(min(h[dis[i].first].first, dis[i].second), T);
			else h[dis[i].first] = mp(dis[i].second, T);
		}
	}
	kk = (LL)kk * inv[a[u]] % mod;
	for(Edge *e = H[u]; e; e = e->next) if(!done[e->v]) {
		int v = e->v;
		nsize = mx[0] = size[v];
		getroot(v, root = 0);
		solve(root);
	}
}

void work()
{
	int u, v;
	for(int i = 1; i <= n; i++) scanf("%d", &a[i]);	
	for(int i = 1; i < n; i++) {
		scanf("%d%d", &u, &v);
		addedges(u, v);
		addedges(v, u);
	}
	ans = mp(INF, INF);
	nsize = mx[0] = n;
	getroot(1, root = 0);
	solve(root);
	if(ans == mp(INF, INF)) printf("No solution\n");
	else printf("%d %d\n", ans.first, ans.second);
}

int main()
{
	for(int i = 1; i < mod; i++) inv[i] = powmod(i, mod - 2);
	while(scanf("%d%d", &n, &kk)!=EOF) {
		init();
		work();
	}

	return 0;
}