牛客练习赛60 E.旗鼓相当的对手

链接

点击跳转

题解

考虑暴力 $dfs$，每次遍历完 $u$的一颗新子树 $v$，假设我建立了一张以深度为下表的表， $cnt[u][i]$表示当前已经遍历过的子树中深度为 $i$的点的个数， $sum[u][i]$表示当前已经遍历过的子树深度为 $i$的点的权值和，那么显然新加入一棵子树会给 $ans[u]$造成贡献：
$\sum_{i=1}^n ( cnt[u][i] \times sum[v][K-i] + sum[u][i] \times cnt[v][K-i] )$

暴力显然会超时

但是观察到这个过程其实是在不断合并子树的，而且计算答案的复杂度是子树大小规模的

那么我就可以直接启发式合并，每次把小的加入大的，可能哈希表会少一个 $log$，但为了省事，我用了 $map$

总复杂度 $O(n\log^2n)$

代码

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define iinf 0x3f3f3f3f
#define linf (1ll<<60)
#define eps 1e-8
#define maxn 200010
#define maxe 200010
#define cl(x) memset(x,0,sizeof(x))
#define rep(i,a,b) for(i=a;i<=b;i++)
#define drep(i,a,b) for(i=a;i>=b;i--)
#define em(x) emplace(x)
#define emb(x) emplace_back(x)
#define emf(x) emplace_front(x)
#define fi first
#define se second
#define de(x) cerr<<#x<<" = "<<x<<endl
using namespace std;
using namespace __gnu_pbds;
typedef long long ll;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
ll read(ll x=0)
{
    ll c, f(1);
    for(c=getchar();!isdigit(c);c=getchar())if(c=='-')f=-f;
    for(;isdigit(c);c=getchar())x=x*10+c-0x30;
    return f*x;
}
ll pool_tot;
ll cnt[maxn], sum[maxn], ans[maxn], n, w[maxn], K;
map<ll,ll> pool[maxn];
ll create()
{
    pool[++pool_tot].clear();
    return pool_tot;
}
ll merge(ll a, ll b)
{
    if(!a or !b)return a+b;
    if(pool[a].size()>pool[b].size())swap(a,b);
    for(auto pr:pool[a])pool[b][pr.fi]+=pr.se;
    return b;
}
ll calc(ll a, ll b, ll S)
{
    if(!a or !b)return 0;
    if(pool[a].size() > pool[b].size())swap(a,b);
    ll ans = 0;
    for(auto pr:pool[a])
    {
        if(pool[b].find(S-pr.fi)!=pool[b].end())
            ans += pr.se*pool[b][S-pr.fi];
    }
    return ans;
}
struct Graph
{
    int etot, head[maxn], to[maxe], next[maxe], w[maxe];
    void clear(int N)
    {
        for(int i=1;i<=N;i++)head[i]=0;
        etot=0;
    }
    void adde(int a, int b, int c=0){to[++etot]=b;w[etot]=c;next[etot]=head[a];head[a]=etot;}
    #define forp(_,__) for(auto p=__.head[_];p;p=__.next[p])
}G;
void dfs(ll u, ll fa, ll depth)
{
    forp(u,G)
    {
        ll v = G.to[p];
        if(v==fa)continue;
        dfs(v,u,depth+1);
        ans[u] += calc(cnt[u],sum[v],K+2*depth);
        ans[u] += calc(cnt[v],sum[u],K+2*depth);
        cnt[u] = merge(cnt[u],cnt[v]);
        sum[u] = merge(sum[u],sum[v]);
    }
    
    auto t = create();
    pool[t][depth] = 1;
    cnt[u] = merge( t, cnt[u] );
    
    t = create();
    pool[t][depth] = w[u];
    sum[u] = merge( t, sum[u] );
}
int main()
{
    ll i, u, v;
    n = read(), K=read();
    rep(i,1,n)w[i] = read();
    rep(i,1,n-1)
    {
        u = read();
        v = read();
        G.adde(u,v);
        G.adde(v,u);
    }
    dfs(1,0,1);
    rep(i,1,n)printf("%lld ",ans[i]);
    return 0;
}