JZOJ4888【NOIP2016提高A组集训第14场11.12】最近公共祖先

Description

YJC最近在学习树的有关知识。今天，他遇到了这么一个概念：最近公共祖先。对于有根树T的两个结点u、v，最近公共祖先LCA(T,u,v)表示一个结点x，满足x是u、v的祖先且x的深度尽可能大。YJC很聪明，他很快就学会了如何求最近公共祖先。他现在想寻找最近公共祖先有什么性质，于是他提出了这样的一个问题：n层的满k叉树T，求对于每一对(i,j)(1≤i,j≤T的点数)，LCA(T,i,j)的深度的和是多少。这个数字n层的满k叉树指一棵带标号的有根树，深度为i($0≤i<n$)的点有k^i个，所有深度≠n-1的点都有k个孩子。YJC发现他不会做了，于是他来问你这个问题的答案。这个答案可能很大，你只需要告诉他答案%998244353的值就可以了。

Solution

我们设f[i]表示i层k叉数的数量，则$f[i]={k^i-1\over{k-1}}$那么我们发现$ans=\sum_{i=0}^{n-1}{(f[n-i+1]^2-k*f[n-i]^2)*i*k^i}$,代入得$ans=\sum_{i=0}^{n-1}{({k^{n-i+1}-1\over{k-1}}^2-k*{k^{n-i}-1\over{k-1}}^2)*i*k^i}$化简一下就得到$ans=\sum^{n-1}_{i=0}{{i*k^i-i*k^{2*n-i+1}}\over{1-k}}$,$ans={\sum^{n-1}_{i=0}{i*k^i-i*k^{2*n-i+1}}\over{1-k}}$我们知道$\sum^{n-1}_{i=0}{i*k^i}=n*\sum^{n-1}_{i=0}{k^i}-\sum_{i=1}^{n}{\sum_{j=0}^{i-1}{k^j}}$,我们将$i*k^{2*n-i+1}$视为$i*k^{2*n+1}*({1\over{k}})^{i}$所以将式子代入得$k^{2n}-k-(2n-1)k^n(k-1)\over{(k-1)^3}$。所以算法复杂度为O($logN$).

Code

#include<iostream>
#include<cmath>
#include<cstring>
#include<cstdio>
#include<algorithm>
#define ll long long
using namespace std;
const ll maxn=1000005,mo=998244353;
ll f[maxn],n,i,t,j,k,l,m,ans,x;
ll sqr(ll x){
    return x*x%mo;
}
int main(){
    freopen("lca.in","r",stdin);freopen("lca.out","w",stdout);
    scanf("%lld%lld",&n,&k);
    f[1]=1;ans=0;n--;
    for (i=2;i<=n+1;i++)
        f[i]=(f[i-1]*k+1)%mo;
    x=1;
    for (i=0;i<=n;i++){
        t=x*((sqr(f[n-i+1])-k*sqr(f[n-i])%mo+mo)%mo)%mo;
        ans=(ans+t*i%mo)%mo;
        x=x*k%mo;
    }
    printf("%lld\n",ans);
}