CodeForces - 1109D Sasha and Interesting Fact from Graph Theory

题意：

给定 $n$ 个带标结点，边权范围 $[1,m]$，再给定两点 $a,b$，求所有可能的 $a,b$ 路径上边权和为 $m$ 的树的种数。$(n,m\le 10^6)$

链接：

https://codeforces.com/contest/1109/problem/D

解题思路：

设 $a,b$ 路径上共有 $s$ 个点，从剩余 $n-2$ 个点中选出 $s$ 排列在 $ab$ 路径上的方案数为 $A_{n-2}^{s-2}$，将 $s-1$ 条边权赋值后和为 $m$ 的方案数为 $C_{m-1}^{s-2}$ （隔板法，$m$ 个球 $m-1$ 个空，插入 $s-2$ 个板，得到 $s-1$ 个划分块），剩余边的边权方案数为 $m^{n-s}$，剩余点生成的森林可以任意拼接在 $ab$ 路径上，由 $Cayle{y}^{\prime }sformula$ 可知，$n$ 个带标结点生成的具有 $s$ 棵树的森林且 $s$ 个指定结点分属不同分量的方案有 $s{n}^{n-s-1}$ 种。综上，
$$ans=A_{n-2}^{s-2}\ast C_{m-1}^{s-2}\ast m^{n-s}\ast s{n}^{n-s-1},s\in [2,n]$$

参考代码：

#include<bits/stdc++.h>
 
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
#define sz(a) ((int)a.size())
#define pb push_back
#define lson (rt << 1)
#define rson (rt << 1 | 1)
#define gmid (l + r >> 1)
const int maxn = 1e6 + 5;
const int inf = 0x3f3f3f3f;
const int mod = 1e9 + 7;

ll inv[maxn], fac[maxn], finv[maxn], pwn[maxn], pwm[maxn];
int n, m, a, b;

ll comb(int n, int m){

    if(m > n) return 0;
    return fac[n] * finv[m] % mod * finv[n - m] % mod;
}

int main(){

    ios::sync_with_stdio(0); cin.tie(0);
    cin >> n >> m >> a >> b;
    inv[1] = fac[0] = finv[0] = pwn[0] = pwm[0] = 1;
    int lim = max(n, m);
    for(int i = 1; i <= lim; ++i){

        inv[i] = i == 1 ? 1 : (mod - mod / i) * inv[mod % i] % mod;
        fac[i] = fac[i - 1] * i % mod;
        finv[i] = finv[i - 1] * inv[i] % mod;
        pwn[i] = pwn[i - 1] * n % mod;
        pwm[i] = pwm[i - 1] * m % mod;
    }
    ll ret = 0;
    for(int i = 2; i <= n; ++i){

        ll tmp = comb(n - 2, i - 2) * fac[i - 2] % mod * comb(m - 1, i - 2) % mod * i % mod * (i == n ? inv[i] : pwn[n - i - 1]) % mod * pwm[n - i] % mod;
        ret = (ret + tmp) % mod;
    }
    cout << ret << endl;
    return 0;
}