【HDU】4812 D Tree 点分治

传送门：【HDU】4812 D Tree

题目分析：点分治搞之。乘积等于K的路径。

首先我们定义一个path[ i ]用以记录从根结点x在子树x内的第 i 条路径的值（乘积）。然后每次我们搞完当前重心rt的一棵子树以后，我们用判断K*逆元[ path[ i ] * val[ rt ] %MOD ] % MOD 是否存在来确定乘积为K的路径是否存在，然后再用这个path[ i ]去更新best[ path[ i ] ]（best[ i ]为取值为 i 的路径的端点最小标号 ）。

求逆元可以用p^(mod - 2 ) % mod来，毕竟互素随便搞？

代码如下：

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std ;

typedef long long LL ;

#pragma comment(linker,"/STACK:102400000,102400000")

#define travel( e , H , u ) for ( Edge* e = H[u] ; e ; e = e -> next )
#define rep( i , a , b ) for ( int i = ( a ) ; i <  ( b ) ; ++ i )
#define rev( i , a , b ) for ( int i = ( a ) ; i >= ( b ) ; -- i )
#define FOR( i , a , b ) for ( int i = ( a ) ; i <= ( b ) ; ++ i )
#define clr( a , x ) memset ( a , x , sizeof a )
#define cpy( a , x ) memcpy ( a , x , sizeof a )

const int MAXN = 100005 ;
const int MAXE = 200005 ;
const int MOD = 1000003 ;
const int INF = 0x3f3f3f3f ;

struct Edge {
	int v ;
	Edge* next ;
} E[MAXE] , *H[MAXN] , *edge ;

struct Node {
	int c , idx ;
	Node () {}
	Node ( int idx , int c ) : idx ( idx ) , c ( c ) {}
} path[MAXN] ;

int flag[MOD] ;
int ni[MOD] ;
int best[MOD] ;
bool vis[MAXN] ;
int val[MAXN] ;
int siz[MAXN] ;
int num[MAXN] ;
int node_num ;
int ans[2] ;
int n , K ;
int root ;
int cnt ;//path
int dep ;//flag

void clear () {
	dep = 0 ;
	edge = E ;
	num[0] = n ;
	clr ( H , 0 ) ;
	clr ( vis , 0 ) ;
	clr ( flag , 0 ) ;
	ans[0] = ans[1] = INF ;
}

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

int pow ( int a , int b ) {
	LL res = 1 , tmp = a ;
	while ( b ) {
		if ( b & 1 ) res = res * tmp % MOD ;
		tmp = tmp * tmp % MOD ;
		b >>= 1 ;
	}
	return res ;
}

void get_siz ( int u , int fa = 0 ) {
	siz[u] = 1 ;
	travel ( e , H , u ) {
		int v = e -> v ;
		if ( !vis[v] && v != fa ) {
			get_siz ( v , u ) ;
			siz[u] += siz[v] ;
		}
	}
}

void get_root ( int u , int fa = 0 ) {
	num[u] = 0 ;
	travel ( e , H , u ) {
		int v = e -> v ;
		if ( !vis[v] && v != fa ) {
			get_root ( v , u ) ;
			num[u] = max ( num[u] , siz[v] ) ;
		}
	}
	num[u] = max ( num[u] , node_num - siz[u] ) ;
	if ( num[u] < num[root] ) root = u ;
}

void get_path ( int u , int c , int fa = 0 ) {
	path[cnt ++] = Node ( u , c ) ;
	travel ( e , H , u ) {
		int v = e -> v ;
		if ( !vis[v] && v != fa ) {
			get_path ( v , ( LL ) c * val[v] % MOD , u ) ;
		}
	}
}

void update_ans ( int u , int v ) {
	if ( u > v ) swap ( u , v ) ;
	if ( u < ans[0] || u == ans[0] && v < ans[1] ) ans[0] = u , ans[1] = v ;
}

void dfs ( int u ) {
	get_siz ( u ) ;
	node_num = siz[u] ;
	root = 0 ;
	get_root ( u ) ;
	int rt = root ;
	vis[rt] = 1 ;
	travel ( e , H , rt ) {
		int v = e -> v ;
		if ( !vis[v] ) dfs ( v ) ;
	}
	++ dep ;
	travel ( e , H , rt ) {
		int v = e -> v ;
		if ( vis[v] ) continue ;
		cnt = 0 ;
		get_path ( v , val[v] % MOD ) ;
		rep ( i , 0 , cnt ) {
			int tmp = ( LL ) val[rt] * path[i].c % MOD ;
			if ( tmp == K ) update_ans ( rt , path[i].idx ) ;
			int x = ( LL ) K * ni[tmp] % MOD ;
			if ( flag[x] == dep ) update_ans ( best[x] , path[i].idx ) ;
		}
		rep ( i , 0 , cnt ) {
			int c = path[i].c , idx = path[i].idx ;
			if ( flag[c] != dep || best[c] > idx ) {
				best[c] = idx ;
				flag[c] = dep ;
			}
		}
	}
	vis[rt] = 0 ;
}

void solve () {
	int x , y , c ;
	clear () ;
	FOR ( i , 1 , n ) scanf ( "%d" , &val[i] ) ;
	rep ( i , 1 , n ) {
		scanf ( "%d%d" , &x , &y ) ;
		addedge ( x , y ) ;
		addedge ( y , x ) ;
	}
	dfs ( 1 ) ;
	if ( ans[0] != INF ) printf ( "%d %d\n" , ans[0] , ans[1] ) ;
	else printf ( "No solution\n" ) ;
}

int main () {
	rep ( i , 0 , MOD ) ni[i] = pow ( i , MOD - 2 ) ;
	while ( ~scanf ( "%d%d" , &n , &K ) ) solve () ;
	return 0 ;
}