树形 dp

「题解」路径计数机

题目链接:NowCoder 1103 B

题解

思路

我们先枚举 $a$ 和 $b$ ,那么满足条件的 $c$ 和 $d$ 一共有两种情况。

  1. $\operatorname{LCA}(c,d)$ 在 $\operatorname{LCA}(a,b)$ 的子树内。
  2. $\operatorname{LCA}(c,d)$ 不在 $\operatorname{LCA}(a,b)$ 的子树内。

对于第一种情况,只要 $\operatorname{LCA}(c,d)$ 不在 $a$ 到 $b$ 的路径上即可。定义 $f _ i$ 为 $\operatorname{LCA}$ 为 $i$,长度为 $q$ 的路径数。预处理 $f$ 数组即可。可以通过枚举每一条路径计算贡献得到。

对于第二种情况,只要 $c$ 到 $d$ 的路径不经过连接 $\operatorname{LCA}(a,b) $ 和 $\operatorname{LCA}(a,b)$ 的父亲的这条边即可。预处理经过每一条边,长度为 $q$ 的路径即可。可以通过枚举每一条路径,然后树上差分计算贡献得到。

代码的渐进时间复杂度为 $\Theta(n^2)$。

代码

渐进时间复杂度为 $\Theta(n^2)$。

#include<cstdio>
#include<algorithm>
using std::max;
#define reg register
typedef long long ll;
#define getchar() (p1==p2&&(p2=(p1=buf)+fread(buf,1,100000,stdin),p1==p2)?EOF:*p1++)
static char buf[100000],*p1=buf,*p2=buf;
inline ll read(void){
    reg bool f=false;
    reg char ch=getchar();
    reg ll res=0;
    while(ch<'0'||'9'<ch)f|=(ch=='-'),ch=getchar();
    while('0'<=ch&&ch<='9')res=10*res+ch-'0',ch=getchar();
    return f?-res:res;
}

const int MAXN=3000+5;
const int MAXLOG2N=12+1;

ll n,p,q;
ll cnt,head[MAXN],to[MAXN<<1],Next[MAXN<<1];
ll f[MAXN][MAXN],g[MAXN][MAXN];
ll fp[MAXN],fq[MAXN],gp[MAXN],gq[MAXN];

inline void Read(void);
inline void Work(void);
inline void Add_Edge(reg int,reg int);
inline void DFS1(reg int,reg int);
inline void DFS2(reg int,reg int);

int main(void){
    Read();
    Work();
    return 0;
}

inline void Read(void){
    n=read(),p=read(),q=read();
    for(reg int i=1;i<n;++i){
        static int u,v;
        u=read(),v=read();
        Add_Edge(u,v);
        Add_Edge(v,u);
    }
    return;
}

inline void Work(void){
    DFS1(1,0),DFS2(1,0);
    reg ll sum1=0,sum2=0;
    for(int i=1;i<=n;++i)
        sum1+=fp[i],sum2+=fq[i];
    reg ll ans=sum1*sum2;
    for(int i=1;i<=n;++i)
        ans-=(fp[i]*fq[i]+fp[i]*gq[i]+fq[i]*gp[i]);
    printf("%lld\n",ans<<2);
    return;
}

inline void Add_Edge(reg int u,reg int v){
    Next[++cnt]=head[u];
    to[cnt]=v;
    head[u]=cnt;
    return;
}

inline void DFS1(reg int u,reg int father){
    f[u][0]=1;
    for(reg int i=head[u];i;i=Next[i]){
        reg int v=to[i];
        if(v!=father){
            DFS1(v,u);
            for(reg int j=1;j<=p;++j)
                fp[u]+=f[v][p-j]*f[u][j-1];
            for(reg int j=1;j<=q;++j)
                fq[u]+=f[v][q-j]*f[u][j-1];
            for(reg int j=1;j<=max(p,q);++j)
                f[u][j]+=f[v][j-1];
        }
    }
    return;
}

inline void DFS2(reg int u,reg int father){
    for(reg int i=1;i<=p;++i)
        gp[u]+=f[u][p-i]*g[u][i];
    for(reg int i=1;i<=q;++i)
        gq[u]+=f[u][q-i]*g[u][i];
    for(reg int i=head[u];i;i=Next[i]){
        reg int v=to[i];
        if(v!=father){
            g[v][1]=1;
            for(reg int j=2;j<=max(p,q);++j)
                g[v][j]+=g[u][j-1]+f[u][j-1]-f[v][j-2];
            DFS2(v,u);
        }
    }
    return;
}