hdu 4812 D Tree 树的点分治

题意 : 给你一棵树 ， 每个点有权值 ， 问是否有一条路径的乘积是K 。

思路 : 基于点的分治 ， 对于一个重心 ，我们枚举以重心为根路径，计算出每个点从路径到根的乘积 ， 然后对于这个乘积 Val， 我们通过hash查找 K * inv[Val] 是否存在 。hash[Val] 我存的到当前重心路径乘积为Val中下标最小的点，如果是INF则表示不存在，为了防止两个端点都在同一子树，要先访问完一遍子树，然后再更新hash。

#include <stdio.h>
#include <string.h>
#include <algorithm>
using namespace std;

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

typedef __int64 LL ;

#define INF 0x3f3f3f3f
#define MAXN 100005
#define MOD 1000003
int ans[2] ;
int n , m , k ;
bool vis[MAXN] ;
int val[MAXN] ;
int inv[MOD] ;
int pi[MOD] ;

struct Tree{
    struct Edge{
        int to , nex ;
        Edge(){}
        Edge( int _to , int _nex ) {
            to = _to ;
            nex = _nex ;
        }
    }edge[MAXN*2] ;

    int head[MAXN] , Index ;

    void init(){
        memset( head , -1 , sizeof(head) ) ;
        Index = 0 ;
    }

    void add( int from , int to ){
        edge[Index] = Edge( to , head[from] ) ;
        head[from] = Index ++ ;
    }
}tree;

int sum[MAXN] , Max[MAXN] ;

void dfs1( int u , int p ) {
    sum[u] = 1 ;
    Max[u] = 0 ;
    for( int i = tree.head[u] ; ~i ; i = tree.edge[i].nex ) {
        int ch = tree.edge[i].to ;
        if( ch == p || vis[ch] ) continue ;
        dfs1( ch , u  ) ;
        sum[u] += sum[ch] ;
        Max[u] = max( Max[u] , sum[ch] ) ; 
    }
}

int Min , Mini ;

void dfs2( int u , int p , int ss ) {
    int MM = max( Max[u] , ss - sum[u] ) ;
    if( MM < Min ) {
        Min = MM ; Mini = u ;
    }
    for( int i = tree.head[u] ; ~i ; i = tree.edge[i].nex ) {
        int ch = tree.edge[i].to ;
        if( ch == p || vis[ch] ) continue ;
        dfs2( ch , u , ss ) ;
    }
}

int getRoot( int u ) {
    dfs1( u , -1 ) ;
    Mini = -1 ;
    Min = 999999999 ;
    dfs2( u , -1 , sum[u] ) ;
    return Mini ;
}

void dfs3( int u , int p , int Val ) {
    Val = ( (LL)Val * val[u] ) % MOD ;
    int key = ( (LL)k * inv[Val])%MOD ;
    if( pi[ key ] != INF ) {
        int tmp[2] ;
        tmp[0] = u ;
        tmp[1] = pi[ key ] ;
        sort( tmp , tmp + 2 ) ;
        if( (tmp[0] < ans[0]) || ( tmp[0] == ans[0] && tmp[1] < ans[1] )) {
            ans[0] = tmp[0] ;
            ans[1] = tmp[1] ;
        }
    }
    for( int i = tree.head[u] ; ~i ; i = tree.edge[i].nex ) {
        int ch = tree.edge[i].to ;
        if( ch == p || vis[ch] ) continue ;
        dfs3( ch , u , Val  ) ;
    }
 }

void dfs4( int u , int p , int Val ) {
    Val = ( (LL)Val * val[u] ) % MOD ;
    pi[Val] = min( pi[Val] , u ) ;

    for( int i = tree.head[u] ; ~i ; i = tree.edge[i].nex ) {
        int ch = tree.edge[i].to ;
        if( ch == p || vis[ch] ) continue ;
        dfs4( ch , u , Val  ) ;
    }
}

void dfs5( int u , int p , int Val ) {
    Val = ( (LL)Val * val[u] ) % MOD ;
    pi[Val] = INF ;
    for( int i = tree.head[u] ; ~i ; i = tree.edge[i].nex ) {
        int ch = tree.edge[i].to ;
        if( ch == p || vis[ch] ) continue ;
        dfs5( ch , u , Val  ) ;
    }
}

void solve( int u ) {
    // 找重心 
    int root = getRoot( u ) ;
    pi[val[root]] = root ;
    for( int i = tree.head[root] ; ~ i ; i = tree.edge[i].nex ) {
        // 计算各个顶点到根节点的π
        int ch = tree.edge[i].to ;
        if( vis[ch] ) continue ;
        // 查询是否存在路径
        dfs3( ch , root , 1 ) ;
        // 在hash中添加路径
        dfs4( ch , root , val[root] ) ;
    }
    for( int i = tree.head[root] ; ~i ; i = tree.edge[i].nex ) {
        int ch = tree.edge[i].to ;
        if( vis[ch] ) continue ;
        dfs5( ch , root , val[root] ) ;
    }
    pi[val[root]] = INF ;
    vis[root] = true ;
    for( int i = tree.head[root] ; ~i ; i = tree.edge[i].nex ) {
        int ch = tree.edge[i].to ;
        if( vis[ch] ) continue ;
        solve( ch ) ;
    }
}

LL power( int a , int k ) {
    if( k == 0 ) return 1 ;
    else if( k == 1 ) return a ;
    LL tmp = power( a , k >> 1 ) ;
    if( k & 1 ) 
        return ( ( tmp * tmp ) % MOD * a ) % MOD ;
    return ( tmp * tmp ) % MOD ;
}

int main(){

    for( int i = 0 ; i < MOD ; i ++ ) {
        inv[i] = power( i , MOD - 2 ) ;
    }

    memset( pi  , INF , sizeof(pi) ) ;

    while( scanf( "%d%d" , &n , &k ) != EOF ) {
        tree.init() ;
        for( int i = 1 ; i <= n ; i ++ ) {
            scanf( "%d" , &val[i] ) ;
        }
        for( int i = 1 ; i < n ; i ++ ) {
            int u , v ;
            scanf( "%d%d" , &u , &v ) ;
            tree.add( u , v ) ;
            tree.add( v , u ) ;
        }
        ans[0] = INF ;
        ans[1] = INF ;
        memset( vis , false , sizeof(vis) ) ;
        solve( 1 ) ;
        if( ans[0] == INF ) {
            puts( "No solution") ;
        }else{
            printf( "%d %d\n" , ans[0] , ans[1] ) ;
        }
    }
    return 0 ;
}