牛客练习赛96

比赛链接

牛客练习赛96

小y的树

题目描述

求一颗\(n\)层的满\(k\)叉树，求任意两点之间距离和等于多少，答案对\(10^9+7\)取模

树上两点距离： 沿着最短路径从u走到v经过的边的数量

\(n\)层的满\(k\)叉树：叶子节点到根的距离是\(n-1\),每个点都有\(k\)个儿子

输入描述:

一行两个正整数代表\(n,k\)

\(2\le n,k\le 10^6\)

输出描述:

输出一行一个数代表答案

输入

3 3

输出

216

解题思路

枚举

考虑每条边的贡献，对于同一层的边的贡献肯定是相同的，假设共有 \(n\) 层边，总的节点数为 \(all=1+k+k^2+\dots +k^n\)，对于第 \(i\) 层的每一条边，下面有 \(now=1+k+\dots +k^{n-i}\) 个点，上面共有 \(all-now\) 个点，则该层的贡献即为 \(k_i\times now\times (all-now)\)

• 时间复杂度：\(O(n)\)

代码

// Problem: 小y的树
// Contest: NowCoder
// URL: https://ac.nowcoder.com/acm/contest/11186/B
// Memory Limit: 524288 MB
// Time Limit: 2000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

// %%%Skyqwq
#include <bits/stdc++.h>
 
#define int long long
#define help {cin.tie(NULL); cout.tie(NULL);}
#define pb push_back
#define fi first
#define se second
#define mkp make_pair
using namespace std;
 
typedef long long LL;
typedef pair<int, int> PII;
typedef pair<LL, LL> PLL;
 
template <typename T> bool chkMax(T &x, T y) { return (y > x) ? x = y, 1 : 0; }
template <typename T> bool chkMin(T &x, T y) { return (y < x) ? x = y, 1 : 0; }
 
template <typename T> void inline read(T &x) {
    int f = 1; x = 0; char s = getchar();
    while (s < '0' || s > '9') { if (s == '-') f = -1; s = getchar(); }
    while (s <= '9' && s >= '0') x = x * 10 + (s ^ 48), s = getchar();
    x *= f;
}

const int N=2e6+5,mod=1e9+7;
int n,k;
int s[N],a[N];
signed main()
{
    cin>>n>>k;
    n--;
    a[0]=1,s[0]=1;
    for(int i=1;i<=n;i++)
    {
    	a[i]=1ll*a[i-1]*k%mod;
    	s[i]=(1ll*a[i]+s[i-1])%mod;
    }
    int res=0;
    for(int i=1;i<=n;i++)
    	res=(res+(s[n]-s[n-i]+mod)%mod*s[n-i]%mod*a[i]%mod)%mod;
    cout<<res;
    return 0;
}