【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%的算法：

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

k2n−k−(2n−1)∗kn∗(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);
}