[ABC337G] Tree Inversion 题解

FFTotoro · 2024-01-26 20:34:50 · 题解

题解区似乎都用的 BIT（树状数组），这里供一个使用 pbds 平衡树的神奇做法。

对于一个结点 u，考虑贡献是怎么来的：一种可能是 v 和 w 在其子树内——此时只需统计每个结点的子树内有多少个编号比它小的，可以使用树上启发式合并，合并的东西是平衡树\color{Red}^1，可以借助 gnu_pbds 内置的 tree 实现；还有一种来源是子树外的（分两种情况，v 在 w 子树内和 v 在 w 子树外只需直接一路下传父亲贡献并去掉 u 本身的贡献），是比较显然的换根 DP。细节比较多注意不要写错。

时间复杂度 O(n\log^2n)，可以通过。

关于 __gnu_pbds::tree 的合并\color{Red}^1：选手可能在赛时使用 swap(A,B) 这种写法来合并，但是与 STL（C++11 往后）不同，在 __gnu_pbds::tree 中这样写的时间复杂度是 O(\max\{|A|,|B|\}) 的，然而 A.swap(B) 却是 O(1) 的，需要注意一下。

#include<bits/stdc++.h>
#include<bits/extc++.h>
#define int long long
using namespace std;
using namespace __gnu_pbds;
typedef tree<int,null_type,less<>,rb_tree_tag,tree_order_statistics_node_update> ordered_set;
main(){
  ios::sync_with_stdio(false);
  int n; cin>>n;
  vector<vector<int> > g(n);
  for(int i=1;i<n;i++){
    int u,v; cin>>u>>v;
    g[--u].emplace_back(--v);
    g[v].emplace_back(u);
  }
  vector<int> fl(n),f1(n),f2(n),f3(n);
  vector<ordered_set> s(n);
  function<void(int,int)> dfs=[&](int u,int r){
    int h_id=-1,h=0;
    for(int i:g[u])
      if(i!=r){
        dfs(i,u),f3[u]+=f3[i];
        if(s[i].size()>h)h=s[i].size(),h_id=i;
      } // 找重儿子
    if(~h_id)s[u].swap(s[h_id]); // 就这儿要注意
    for(int i:g[u])
      if(i!=r&&i!=h_id)
        for(int j:s[i])s[u].insert(j); // 暴力合并
    int x=s[u].order_of_key(u);
    f3[u]+=(fl[u]=x),s[u].insert(u);
    f1[u]=s[u].order_of_key(r);
  }; // 预处理及计算子树内的答案
  function<void(int,int)> dfs2=[&](int u,int r){
    if(u)f2[u]=f2[r]+f3[r]-f3[u]+(u-fl[u])-f1[u];
    for(int i:g[u])
      if(i!=r)dfs2(i,u);
  }; // 计算子树外的答案
  dfs(0,0),dfs2(0,0);
  for(int i=0;i<n;i++)
    cout<<f2[i]+f3[i]<<' ';
  cout<<endl;
  return 0;
}