[ABC312G] Avoid Straight Line 题解

· · 题解

前言

纪念自己第二次做出 ABC 的 G！

解法

显然的一个树形 DP。

令 e_{r,i} 为在以 r 为根的树中以 i 为根结点的子树大小，f_{r,i} 为在以 r 为根的树中结点 i 的父亲。先将这棵树视为以 1 为根的有根树，维护出所有 e_{1,i}。

正难则反，考虑有多少个 (i,j,k) 在同一条路径上。枚举三个点中间的那一个 x，剩下两个只可能在它的几个儿子管辖的子树内或在 x 管辖的子树外（当然另外两个点不可能在同一棵子树内或都在 x 管辖的子树外）。

为方便理解，我们将 x 视为根，令以 r 为根的子树中的 x 儿子集合为 \mathrm{son}_{r,x}，这时以 1 为根的树中 x 的子树外的所有结点在以 x 为根的情况下变成了以 f_{1,x} 为根的一棵子树，其大小为 n-e_{1,x}。显然以 x 为中点的路径数为：

\begin{aligned} \sum\limits_{u\in\mathrm{son}_{x,x}}\sum\limits_{v\in\mathrm{son}_{x,x}\land u\ne v}e_{x,u}e_{x,v}=\dfrac{\left(n-e_{1,x}+\sum\limits_{u\in\mathrm{son}_{1,x}} e_{1,u}\right)^2-\left(n-e_{1,x}\right)^2-\sum\limits_{u\in\mathrm{son}_{1,x}}e_{1,u}^2}{2} \end{aligned}

这里 \land 表示合取。

把所有点的贡献加起来，记为 s，再用 \mathrm{C}_n^3=\dfrac{n(n-1)(n-2)}{6} 减去 s 即为答案。

放代码：

#include<bits/stdc++.h>
#define int long long
using namespace std;
main(){
  ios::sync_with_stdio(false);
  cin.tie(0); cout.tie(0);
  int n,s=0; cin>>n;
  vector<int> e(n),f(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);
  }
  function<void(int,int)> dfs=[&](int u,int f){
    e[u]=1;
    for(int i:g[u])
      if(i!=f)dfs(i,u),e[u]+=e[i];
  };
  dfs(0,0);
  function<void(int,int)> dfs2=[&](int u,int f){
    int c=0,c2=0;
    for(int i:g[u])
      if(i==f)c+=n-e[u],c2+=(n-e[u])*(n-e[u]);
      else dfs2(i,u),c+=e[i],c2+=e[i]*e[i];
    s+=c*c-c2>>1;
  };
  dfs2(0,0);
  cout<<n*(n-1)*(n-2)/6-s<<endl;
  return 0;
}