P4427 [BJOI2018]求和(LCA+字首和)

小菜雞加油發表於2020-11-21

https://www.luogu.com.cn/problem/P4427


 

題目描述

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

輸入格式

第一行包含一個正整數nn,表示樹的節點數。

之後n-1n−1 行每行兩個空格隔開的正整數i, ji,j,表示樹上的一條連線點ii 和點jj 的邊。

之後一行一個正整數mm,表示詢問的數量。

之後每行三個空格隔開的正整數i, j, ki,j,k,表示詢問從點ii 到點jj 的路徑上所有節點深度的kk 次方和。由於這個結果可能非常大,輸出其對998244353998244353 取模的結果。

樹的節點從11 開始標號,其中11 號節點為樹的根。

輸出格式

對於每組資料輸出一行一個正整數表示取模後的結果。


思路:暴力統計50次,記錄字首和,然後類似樹上兩點距離,用lca來優化。lca用樹鏈可求

#include<iostream>
#include<vector>
#include<queue>
#include<cstring>
#include<cmath>
#include<map>
#include<set>
#include<cstdio>
#include<algorithm>
#define debug(a) cout<<#a<<"="<<a<<endl;
using namespace std;
const int maxn=3e5+1000;
typedef long long LL;
const LL mod=998244353;
LL ksm(LL a,LL k){LL res=1;while(k){if(k&1) res=res*a%mod;a=a*a%mod;k>>=1;}return res%mod;}
vector<LL>g[maxn];
LL dep[maxn],fa[maxn],siz[maxn],son[maxn];
LL top[maxn];
LL sum[maxn][51];
void predfs(LL u,LL father){
    siz[u]=1;fa[u]=father;
    for(LL i=0;i<g[u].size();i++){
        LL v=g[u][i];
        if(v==father) continue;
        dep[v]=dep[u]+1;
        for(LL k=1;k<=50;k++)
        sum[v][k]=(sum[u][k]%mod+ksm(dep[v],k)%mod)%mod;
        predfs(v,u);
        siz[u]+=siz[v];
        if(siz[v]>siz[son[u]]){
            son[u]=v;
        }
    }
}
void dfs(LL u,LL topx){
    top[u]=topx;
    if(!son[u]) return;
    dfs(son[u],topx);
    for(LL i=0;i<g[u].size();i++){
        LL v=g[u][i];
        if(v==fa[u]||v==son[u]) continue;
        dfs(v,v);
    }
}
LL lca(LL u,LL v){
    while(top[u]!=top[v]){
        if(dep[top[u]]<dep[top[v]]) swap(u,v);
        u=fa[top[u]];
    }
    return dep[u]<dep[v]?u:v;
}
int main(void)
{
  cin.tie(0);std::ios::sync_with_stdio(false);
  LL n;cin>>n;
  for(LL i=1;i<n;i++){
    LL u,v;cin>>u>>v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dep[0]=-1;dep[1]=0;
  predfs(1,0);
  dfs(1,1);
  ///cout<<"fuck"<<endl;
  LL m;cin>>m;
  while(m--)
  {
      LL u,v,k;cin>>u>>v>>k;
      LL LCA=lca(u,v);
      cout<<( (sum[u][k]%mod+sum[v][k]%mod)%mod-(sum[LCA][k]%mod+sum[fa[LCA]][k]%mod)%mod+mod)%mod<<endl;
  }
return 0;
}

 

相關文章