HDU 4670 Cube number on a tree ( 树的点分治 )

Posted 2020-09-10 yangykaifa

题意 : 给你一棵树 。 树的每一个结点都有一个权值 。 问你有多少条路径权值的乘积是一个全然立方数 。

题目中给了你 K 个素数 ( K <= 30 ) , 全部权值都能分解成这k个素数

思路 : 一个全然立方数的素因子个数都是三的倍数 ， 所以我们仅仅要求各个素数的个数即可了 ， 而且我们仅仅关心个数对三的余数

所以我们能够用一个 长整形来表示每一个结点到根的各个素因子的个数( 三进制压缩 ) 。

只是由于用位运算会快一点 ， 所以我用了四进制。即每两位表示一个素因子的个数 。

中间合并的时候计算刚好相加之后变成0的数，然后用map统计就能够了

注意 : 题目中 pi 的值 ， 0 是取不到的(  纠结了好久该怎么处理 = = )

#include <stdio.h>
#include <string.h>
#include <map>
#include <algorithm>
#pragma comment(linker, "/STACK:102400000,102400000")
using namespace std;

#define MAXN 50005
#define INF 0x3f3f3f3f
__int64 prime[35] ;
int n , k ;

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

    int head[MAXN] ;  
    int 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 ;

inline int get( __int64 val , int pos ) {
    return ( val >> ( pos << 1 ) ) & 3 ;
}

inline void set( __int64 & val , int pos , int k ) {
    val |= ( (__int64)k << ( pos << 1 ) ) ;
}

__int64 add( __int64 a , __int64 b ) {
    __int64 ans = 0 ;
    for( int i = 0 ; i < k ; i ++ ) {
        int aa = get( a , i ) ;
        int bb = get( b , i ) ;
        aa += bb ;
        aa %= 3 ; ;
        set( ans , i , aa ) ;
    }
    return ans ;
}

__int64 fan( __int64 a  ) {
    __int64 ans = 0 ;
    for( int i = 0 ; i < k ; i ++ ) {
        int tmp = get( a , i ) ;
        tmp =  3 - tmp ;
        set( ans , i , tmp % 3 ) ;
    }
    return ans ;
}

bool vis[MAXN] ;
__int64 val[MAXN] ;
__int64 ans ;

// 第i结点为根的子树的最大个数  
int dp[MAXN] ;  
// 第i结点为根的子树的大小  
int sum[MAXN] ;  

int Sum ;  
int Min , Minid ;  

void dfs1( int u , int f ) {  
    sum[u] = 1 ;  
    dp[u] = 0 ;  
    for( int i = tree.head[u] ; ~i ; i = tree.edge[i].nex ) {  
        int ch = tree.edge[i].to ;  
        if( ch == f || vis[ch] ) continue ;  
        dfs1( ch , u  ) ;  
        sum[u] += sum[ch] ;  
        dp[u] = max( dp[u] , sum[ch] ) ;  
    }  
}  
  
void dfs2( int u , int f ) {  
    int M = max( dp[u] , Sum - sum[u] ) ;  
    if( M < Min ) {  
        Min = M ;  
        Minid = u ;  
    }  
    for( int i = tree.head[u] ; ~i ; i = tree.edge[i].nex ) {  
        int ch = tree.edge[i].to ;  
        if( ch == f || vis[ch] ) continue ;  
          
        dfs2( ch , u ) ;  
    }  
}  
  
// 返回树的重心  
int getRoot( int u ) {  
    // 第一遍dfs求出全部点的子树大小。和以孩子为跟的子树个数的最大值  
    dfs1( u , 0 ) ;  
    Sum = sum[u] ;  
    Min = 0x3f3f3f3f ;  
    Minid = -1 ;  
    // 第二次dfs求重心  
    /* 事实上假设整棵树求重心的话，我们一边dfs即可。由于这个是部分树，我们要dfs来确定哪些结点在当前树中 */  
    dfs2( u , 0 ) ;  
    return Minid ;  
}  

int tot ;
__int64 Val[MAXN] ;

void getVal( int u , int fa , __int64 vv ) {
    Val[tot++] = add( vv , val[u] ) ;
    for( int i = tree.head[u] ; ~i ; i = tree.edge[i].nex ) {
        int to = tree.edge[i].to ;
        if( vis[to] == true ) 
            continue ;
        if( to == fa ) continue ;
        getVal( to , u , add( vv , val[u] ) ) ;
    }
}

map<__int64,int> mp ;

__int64 cal( int root , __int64 val , __int64 root_val ) {
    tot = 0 ;
    getVal( root , 0 , val ) ;
    //sort( Val , Val + tot ) ;
    mp.clear() ;
    __int64 ans = 0 ;
    for( int i = 0 ; i < tot ; i ++ ) {
        ans += mp[ add( root_val , fan(Val[i]) ) ] ;
        mp[Val[i]] ++ ;
    }
    return ans ;
}

void solve( int u ) {  
    // 得到当前子树的重心  
    int root = getRoot( u ) ;  
    // 计算以root为根的结果  
    ans += cal( root , 0 , val[root] ) ;  
    if( val[root] == 0 ) ans ++ ;
    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 ;  
        ans -= cal( ch , val[root] , val[root] ) ;  
        solve( ch ) ;     
    }  
}

int main(){
    while( scanf( "%d" , &n ) != EOF )  {
        scanf( "%d" , &k ) ;
        for( int i = 0 ; i < k ; i ++ ) {
            scanf( "%d" , &prime[i] ) ;
        }
        tree.init() ;
        memset( vis , false , sizeof(vis) );
        memset( val , 0 , sizeof(val) ) ;
        for( int i = 1 ; i <= n ; i ++ ) {
            __int64 tmp ;
            scanf( "%I64d" , &tmp ) ;
            for( int j = 0 ; j < k ; j ++ ) {
                int cnt = 0 ;
                while( tmp % prime[j] == 0 ) {
                    cnt ++ ;
                    tmp /= prime[j] ;
                }
                cnt %= 3 ;
                set( val[i] , j , cnt ) ;
            }
        }
        for( int i = 0 ; i < n - 1 ; i ++ ) {
            int u , v ;
            scanf( "%d%d" , &u , &v ) ;
            tree.add( u , v ) ;
            tree.add( v , u ) ;
        }
        ans = 0 ;
        solve( 1 ) ;
        printf( "%I64d\n" , ans ) ;
    }
    return 0 ;
}