题解：CF954H Path Counting

题目传送门

题意

给出一颗树，树的每一层的节点的儿子数目相同。求出当 \(k\) 等于 \(1 \sim 2\times n-2\) 时，树上有多少条不同路径满足路径长度等于 \(k\)。

思路

一道非常有意思的推式子题。
首先我们定义：

\[g_i=\prod_{k=1}^{j}a_k \]

\[m(i,j)=\prod_{k=i+1}^{i+j-1}a_k=\frac{g_{i+j-1}}{g_i} \]

\(g_i\) 就是前缀积，\(m(i,j)\)表示以深度为 \(i\) 的点为端点，向下延伸 \(j\) 的长度的路径有多少条。这个很好证明，我每到一个点 \(i\) 都会有 \(a_i\) 的路径可供选择，就给答案乘上 \(a_i\)，所以路径数就是 \(a_i\) 的乘积。

但是我们只考虑了向下延伸的路径，还会有什么路径呢？显然，一条路径有可能是弯折的，也就是说，可以被拆分成两段，两段是由自己的两个不同的儿子延伸上来的。

那么这个路径该怎么算呢。我们假设当前要求的路径长度为 \(len\)，当前枚举到的节点深度为 \(i\)，那么在这一层相同情况的节点共有 \(g_{i-1}\) 个，我们枚举从任意两个儿子中选出两条路径拼一起使得长度为 \(len\)。那么选取任意两个儿子有 \(\frac{(a_i-1)*a_i}{2}\) 种情况，我们再枚举儿子，满足相加等于 \(len\) 的情况有：\(\sum_{j=1}^{len} m(i+1,j-1)\times m(i+1,j+len-1)\)。为什么是 \(i+1\)呢，因为要枚举的是儿子。而为什么长度都要减一呢，因为要减去从当前节点到儿子的路径长度。

最终计算的式子就是：

\[ans_{len}=\sum_{i=1}^{n-1}\frac{(a_i-1)*a_i}{2}\times g_{i-1}\sum_{j=1}^{len} m(i+1,j-1)\times m(i+1,len-j-1) \]

用上前缀积优化可以做到 \(O(n^3)\)。
这里贴上部分代码：

for(int i=n-1;i>=1;i--){
    ll u=i&1,v=u^1;
    sum=sum*a[i]%mod;
    ans[1]=(ans[1]+a[i]*mul[i-1]%mod)%mod;
    for(int j=1;j<n-i+1;j++){
        for(int k=1;k<=j;k++){
            if(j==k){
                ll tmp=(a[i]-1)*a[i]%mod*Mod%mod;
                tmp=tmp*dp[u][j]%mod*dp[u][k]%mod*mul[i-1]%mod;
                ans[j+k]=(ans[j+k]+tmp)%mod;
                }
            else{
                ll tmp=(a[i]-1)*a[i]%mod;
                tmp=tmp*dp[u][j]%mod*dp[u][k]%mod*mul[i-1]%mod;
                ans[j+k]=(ans[j+k]+tmp)%mod;
            }
        }
    }
    dp[v][1]=1;
    if(i==1)break;
    for(int j=2;j<=n-i+1;j++){
        ans[j]=(ans[j]+dp[u][j-1]*a[i]%mod*mul[i-1]%mod)%mod;
		dp[v][j]=dp[u][j-1]*a[i]%mod;
	}
}

但是这肯定是不够的，现在考虑如何优化成 \(O(n^2)\)。

首先对我们刚刚那个式子进行变形：

\[ans_{len}=\sum_{i=1}^{n-1}\frac{(a_i-1)*a_i}{2}\times g_{i-1}\times \frac{1}{g_{i}^2}\sum_{j=1}^{len} g_{i+j-1}\times g_{i+len-j-1} \]

这是把 \(m(i,j)\) 拆成了 \(\frac{g_{i+j-1}}{g_{i}}\) 的形式，并把 \(\frac{1}{g_i}\) 提出来，由于各有一个所以要平方。

接下来我们令 \(dp_{i,j}=\sum_{j=1}^{len} g_{i+j-1}\times g_{i+len-j-1}\)。

观察下面几个等式：

\[dp_{2,i}=g_{i}^2 \]

\[dp_{3,i}=2\times g_i\times g_{i+1} \]

\[dp_{4,i}=g_{i+1}^2+2\times g_i\times g_{i+2}=dp_{2,i+1}+2\times g_i\times g_{i+4-2} \]

综上所述，我们可以总结出规律：

\[dp_{i,j}=dp_{i-2,j+1}+2\times g_j\times g_{i+j-2} \]

那么我们就可以用 \(O(n^2)\) 的复杂度求出 \(dp\) 数组了。

注意一个小地方，由于 \(n\le5000\) 如果给 \(dp\) 数组开 \(n^2\) 的空间的话会爆空间，所以我们可以开一个滚动数组。

完整代码如下：

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=1010101;
const int mod=1e9+7;
const int Mod=5e8+4;
ll n,mul[N],a[N],dp[5][N],ans[N],inv[N];
ll qpow(ll x,ll y){
	ll res=1;
	while(y){
		if(y&1)res=res*x%mod;
		x=x*x%mod;
		y>>=1;
	}
	return res;
}
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    cin>>n;
    mul[0]=1;
    for(int i=1;i<n;i++)cin>>a[i];
    for(int i=1;i<n;i++)mul[i]=mul[i-1]*a[i]%mod,inv[i]=qpow(mul[i],mod-2);
    for(int len=1;len<n;len++)
        for(int i=1;i<=n-len;i++)
            ans[len]=(ans[len]+mul[i+len-1])%mod;
    for(int i=1;i<n;i++)dp[2][i]=mul[i]*mul[i]%mod;
    for(int i=2;i<=2*n-2;i++){
        ll u=i%3,v=u+1;
        v%=3;
        for(int j=1;j<n;j++){
            if(i!=2)dp[u][j]=(dp[v][j+1]+2*mul[j]%mod*mul[i+j-2]%mod)%mod;
            ans[i]=(ans[i]+(a[j]-1)*Mod%mod*mul[j]%mod*inv[j]%mod*inv[j]%mod*dp[u][j]%mod)%mod;
        }
    }      
    for(int i=1;i<=2*n-2;i++)cout<<ans[i]<<" ";
    return 0;
}