P5628 【AFOI-19】面基

不是哥们，这题有什么好容斥的？直接换根啊！

观察到

n\times k

很小，设计

f_{i,j}

表示子树

i

中距

i

为

j

的边的贡献总和。一个 dfs 可以快速解决。

我们应该要解决以每个点为根的答案的

\max

。

对于以

u

为根的树，答案应为

\sum^{k}_{i=0} f_{u,i}

，然后枚举所有

u\to v

，先在

f_u

中减去

f_v

的贡献，同时更新子树大小，然后再在

f_v

中加入

f_u

。

代码很少很清晰。

CPP

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define pb push_back
#define eb emplace_back
#define all(x) (x).begin(),(x).end()
#define fi first
#define se second
#define SZ(x) ((int)(x).size())
typedef vector<int> VI;
typedef long long ll;
typedef pair<int,int> PII;
const ll MOD=1e9+7;
// head
const int N=3e4+5,M=205;
int n,k;
vector<vector<int>> G(N);
int sz[N],f[N][M];
void dfs(int u,int fa)
{
    sz[u]=1;
    for(auto v:G[u]){
        if(v==fa) continue;
        dfs(v,u);
        sz[u]+=sz[v];
        for(int i=1;i<=k;i++) f[u][i]+=f[v][i-1];
        f[u][0]+=sz[v]*(n-sz[v]);
    }
}
int ans;
void dfs2(int u,int fa)
{
    int sum=0;
    for(int i=0;i<=k;i++) sum+=f[u][i];
    ans=max(ans,sum);
    for(auto v:G[u]){
        if(v==fa) continue;
        for(int i=1;i<=k;i++) f[u][i]-=f[v][i-1];
        f[u][0]-=sz[v]*(n-sz[v]);
        sz[u]-=sz[v],sz[v]+=sz[u];
        for(int i=1;i<=k;i++) f[v][i]+=f[u][i-1];
        f[v][0]+=sz[u]*(n-sz[u]);
        dfs2(v,u);
        for(int i=1;i<=k;i++) f[v][i]-=f[u][i-1];
        f[v][0]-=sz[u]*(n-sz[u]);
        sz[v]-=sz[u],sz[u]+=sz[v];
        for(int i=1;i<=k;i++) f[u][i]+=f[v][i-1];
        f[u][0]+=sz[v]*(n-sz[v]);
    }
}
signed main() 
{
    cin.tie(nullptr);
    ios::sync_with_stdio(false);

    cin>>n>>k;
    for(int i=1;i<n;i++){
        int u,v;cin>>u>>v;
        G[u].pb(v);
        G[v].pb(u);
    }
    dfs(1,0);
    dfs2(1,0);
    cout<<ans<<'\n';
}

P5628 【AFOI-19】面基

文章操作

相关推荐

评论

P5628 【AFOI-19】面基