【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的值就可以了。
Input
第一行包含两个整数n和k，表示T是一棵n层的满k叉树。
Output
一行，包含一个整数，表示问题的答案%998244353的值。
Sample Input
3 2
Sample Output
22
Data Constraint
​对于30%的数据，满足2≤n,k≤8；
对于50%的数据，满足2≤n,k≤1000000；
对于100%的数据，满足2≤n,k≤998244351。
Hint
LCA深度为0的点对有31个，深度为1的点对有14个，深度为2的点对有4个，所以答案=31*0+14*1+4*2=22。

The Solution

30%的算法 暴力求lca，直接计算。
100%的算法：

手玩出一大波算式，然后整理一下可得出

$$\dfrac{k^{2n}-k-(2n-1)*k^n*(k-1)}{(k-1)^3}$$

最后就随便打了。。

CODE

#include <cstdio>
#include <iostream>
#include <algorithm>
#include <cmath>
#include <cstring>
#define fo(i,a,b) for (int i=a;i<=b;i++)
#define mo 998244353
#define N 100005

using namespace std;

typedef long long ll;

ll ni,n,k;

ll mi(ll x,int y )
{
    ll z = 1;
    for (;y;x = x * x % mo,y >>= 1) if (y & 1) z = z * x % mo;
    return z;
}


int main()
{
    freopen("lca.in","r",stdin);
    freopen("lca.out","w",stdout);
    scanf("%lld%lld",&n,&k); 
    ni = mi(mi(k-1,3),mo - 2); 
    ll ans = (mi(k,2*n) % mo + mo +mo - k - (2*n-1) * mi(k,n) % mo * (k-1) % mo) % mo;  
    printf("%lld\n",ans * ni % mo);
}