B题的DP解法,比赛的时候脑抽了,想了个DP还挺合理的,但是写代码的时候少算一项,一直过不了样例,赛后调过去了
DP[0]:表示无节点在最后一层的方案数量;
DP[1]:表示有一节点在最后一层的方案数量;
DP[2]:表示有二节点在最后一层的方案数量;
SUM[0]:表示有无节点在最后一层的总长度;
SUM[1]:表示有一节点在最后一层的总长度;
SUM[1]:表示有二节点在最后一层的总长度;
画个图就知道什么意思了

那我们只要考虑每次加一层,DP,和SUM会怎么改变就好了,需要考虑DP,SUM之间的关系,具体可看代码,就可以知道转移方程:
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int mod = 1e9 + 7;
ll dp[5];
ll sum[5];
void run() {
ll n, k;
cin >> n >> k;
dp[0] = 0;
dp[1] = k;
dp[2] = k * (k - 1) / 2 % mod;
sum[0] = 0;
sum[1] = k;
sum[2] = k * (k - 1) % mod;
ll cnt = k;
for (int i = 1; i < n - 1; i++) {
ll _dp[5];
ll _sum[5];
memcpy(_dp, dp, sizeof dp);
memcpy(_sum, sum, sizeof sum);
dp[0] = (_dp[0] + _dp[1] + _dp[2]) % mod;
sum[0] = (_sum[0] + _sum[1] + _sum[2]) % mod;
dp[1] = (_dp[1] * k % mod + 2 * _dp[2] * k % mod + cnt * k % mod) % mod;
sum[1] = (_sum[1] * k % mod + dp[1] + _sum[2] * 2 * k % mod ) % mod;
dp[2] = (_dp[2] * k % mod * k % mod + k * (k - 1) / 2 % mod * cnt) % mod;
sum[2] = (_sum[2] * k % mod * k % mod + dp[2] * 2) % mod;
cnt *= k;
cnt %= mod;
}
cout << (sum[0] + sum[1] + sum[2]) % mod << endl;
}
int main() {
// int t;
// cin >> t;
// while (t--)
run();
return 0;
}