AT_abc337_g Tree Inversion

原题链接

换根 dp，先随便钦定一个根。

记 \(a_x\) 为以 \(x\) 为根的子树中小于 \(x\) 的点的个数，\(g_x\) 为以 \(x\) 为根的子树中小于 \(fa_x\) 的点的个数，也就是 \(x\) 对 \(f_{fa_x}\) 的贡献。

上述两个东西可以直接在按 DFS 序差分加树状数组求出。

void dfs(int x,int fa=0)
{
    dfn[x]=++tot,pos[tot]=x,siz[x]=1;
    bit::change(x,1);
    for(int y:v[x])
    {
        if(y==fa) continue;
        //这里计算 y 对 x 的贡献
        g[y]-=a[x];a[x]-=bit::query(x-1);
        dfs(y,x),siz[x]+=siz[y];
        a[x]+=bit::query(x-1);g[y]+=a[x];
    }
    return ;
}

然后考虑对于一个点 \(p\) 的答案，不在其到根的路径上的点 \(x\) 此时对答案的贡献是 \(a_x\)；而对于在其到根路径上的点 \(x\)，\(y\) 是子树中包含 \(p\) 的 \(x\) 的儿子，\(x\) 的贡献应当是所有小于 \(x\) 的数的个数减去以 \(y\) 为根的子树中小于 \(x\) 的个数即 \((x-1)-g_y\)（对于 \(p\) 本身来说贡献就是 \(p-1\)）。

所以可以先求出 \(\sum a_i\)，然后从根开始再 DFS 一遍，从 \(x\) 往 \(y\) 递归时更改 \(x\) 的贡献即可。

void dp(int x,int sum,int fa=0)
{
    ans[x]=sum-a[x]+(x-1);
    for(int y:v[x])
    {
        if(y==fa) continue;
        dp(y,sum-a[x]+(x-1)-g[y],x);
    }
    return ;
}
signed main()
{
    cin.tie(0),cout.tie(0);
    ios::sync_with_stdio(0);
    cin>>n;
    for(int i=1,x,y;i<n;++i)
        cin>>x>>y,v[x].push_back(y),v[y].push_back(x);
    dfs(1);for(int i=1;i<=n;++i) sum+=a[i];
    dp(1,sum);for(int i=1;i<=n;++i) cout<<ans[i]<<' ';
    cout<<'\n';return 0;
}