color 不同 swap pre pac truct desc con 不同的

Description

master 對樹上的求和非常感興趣。他生成了一棵有根樹，並且希望多次詢問這棵樹上一段路徑上所有節點深度的k 次方和，而且每次的k 可能是不同的。此處節點深度的定義是這個節點到根的路徑上的邊數。他把這個問題交給 了pupil，但pupil 並不會這麽復雜的操作，你能幫他解決嗎？

Input

第一行包含一個正整數n ，表示樹的節點數。 之後n-1 行每行兩個空格隔開的正整數i,j ，表示樹上的一條連接點i 和點j 的邊。 之後一行一個正整數m ，表示詢問的數量。 之後每行三個空格隔開的正整數i,j,k ，表示詢問從點i 到點j 的路徑上所有節點深度的k 次方和。 由於這個結果可能非常大，輸出其對998244353 取模的結果。 樹的節點從1 開始標號，其中1 號節點為樹的根。

Output

對於每組數據輸出一行一個正整數表示取模後的結果。 1≤n,m≤300000,1≤k≤50

Sample Input

5
1 2
1 3
2 4
2 5
2
1 4 5
5 4 45

Sample Output

33
503245989
說明
樣例解釋
以下用d(i) 表示第i 個節點的深度。
對於樣例中的樹，有d(1)=0,d(2)=1,d(3)=1,d(4)=2,d(5)=2。
因此第一個詢問答案為(2^5 + 1^5 + 0^5) mod 998244353 = 33
第二個詢問答案為(2^45 + 1^45 + 2^45) mod 998244353 = 503245989。

Solution

因為這個k很小，所以就是個樹上差分板子

預處理出深度之後弄個樹上前綴和（對於不同的k暴力存就好，因為k最大50）

預處理復雜度是$O(nk)$的

然後每次詢問的答案其實就是兩個點的前綴和減掉他們的$LCA$的前綴和再加上他們的$LCA$的k次方

#include <bits/stdc++.h>

using namespace std ;

#define N 300010
#define mod 998244353
#define inf 0x3f3f3f3f
#define ll long long

int n , m ;
int dep[ N ] , siz[ N ] , top[ N ] , fa[ N ] ;
 
int head[ N ] , cnt ;
ll c[ N ][ 60 ] ; 
struct node {
    int to , nxt ;
} e[ N << 1 ] ;

void ins( int u , int v ) {
    e[ ++ cnt ].to = v ;
    e[ cnt ].nxt = head[ u ] ; 
    head[ u ] = cnt ;
}

void dfs1( int u ) {
    siz[ u ] = 1 ; 
    for( int i = head[ u ] ; i ; i = e[ i ].nxt ) {
        int v = e[ i ].to ;
        if( v == fa[ u ] ) continue ;
        fa[ v ] = u ; 
        dep[ v ] = dep[ u ] + 1 ;
        dfs1( v ) ;
        siz[ u ] += siz[ v ] ;
    }
}

void dfs2( int u , int topf ) {
    top[ u ] = topf ;
    int k = 0 ;
    for( int i = head[ u ] ; i ; i = e[ i ].nxt ) {
        if( e[ i ].to != fa[ u ] && siz[ e[ i ].to ] > siz[ k ] ) 
            k = e[ i ].to ; 
    }
    if( !k ) return ;
    dfs2( k , topf ) ;
    for( int i = head[ u ] ; i ; i = e[ i ].nxt ) {
        if( e[ i ].to != k && e[ i ].to != fa[ u ] ) 
            dfs2( e[ i ].to , e[ i ].to ) ;
    }
}

void dfs3( int u ) {
    ll x = 1 ;
    for( int k = 1 ; k <= 50 ; k ++ ) {
        x = 1ll * x * dep[ u ] % mod ;
        c[ u ][ k ] = ( 1ll * c[ fa[ u ] ][ k ] + x ) % mod ;
    }
    for( int i = head[ u ] ; i ; i = e[ i ].nxt ) {
        if( e[ i ].to == fa[ u ] ) continue ;
        dfs3( e[ i ].to ) ;
    }
}

int lca( int x , int y ) {
    while( top[ x ] != top[ y ] ) {
        if( dep[ top[ x ] ] < dep[ top[ y ] ] ) swap( x , y ) ;
        x = fa[ top[ x ] ] ;
    }
    if( dep[ x ] > dep[ y ] ) swap( x , y ) ;
    return x ;
}

ll power( ll a , ll b ) {
    ll ans = 1 , base = a ;
    while( b ) {
        if( b&1 ) ans = ans * base % mod ;
        base = base * base % mod ;
        b >>= 1 ;
    }
    return ans % mod ;
}

int main() {
    scanf( "%d" , &n ) ;
    for(int i = 1 , u , v ; i < n ; i ++ ) {
        scanf( "%d%d" ,  &u , &v ) ;
        ins( u ,  v ) ; ins( v , u ) ;
    }
    dfs1( 1 ) ; dfs2( 1 , 1 ) ; dfs3( 1 ) ;
    scanf( "%d" , &m ) ;
    while( m -- ) {
        int u , v , k ;
        scanf( "%d%d%d" , &u , &v , &k ) ;
        int l = lca( u , v ) ;
        printf( "%lld\n" , ( c[ u ][ k ] % mod + c[ v ][ k ] % mod - 2 * c[ l ][ k ] % mod + power( dep[ l ] , k ) + mod ) % mod ) ;
    }
    return 0 ; 
} 

BZOJ5293: [Bjoi2018]求和 樹上差分