AtCoder Beginner Contest 312F問題

\(\def \set #1#2{\{ #1 \ \vert \ #2 \}}\) \(\def \order #1{\lvert {#1} \rvert}\)

問題概要

\(N\) 頂点の tree \(G\) が与えられる． \begin{align} U &:= \set{(i,j,k) \in V(G)}{ i,j,k \text{は相異なる} }, \\ L &:= \set{(i,j,k) \in U}{i,j,k \text{は一直線状に並ぶ}} \end{align} とおく． ここで，\(i,j,k\) が一直線上に並ぶとは， 次のいずれかが成り立つことである．

• \(k \in V(mini\text{-}path_{i,j})\),
• \(i \in V(mini\text{-}path_{j,k})\),
• \(j \in V(mini\text{-}path_{k,i})\).

解法

答えは \[ \begin{matrix} ans &=& \order{U} - \order{L} \\ &=& {}_{N} C_{3} - \sum_{i,j \in V(G) \\ i \neq j} (dis_{i,j} - 1)\\ &=& {}_{N} C_{3} - \sum_{e \in E(G)}\order{ \set{ (i,j) \in V(G)^{2} }{ i \neq j \text{かつ} e \in E(mini\text{-}path_{i,j})} } + {}_{N}C_{2} \end{matrix} \] となる．

頂点の組 \((i,j)\) を動かす代わりに， 辺 \(e = (i,j)\) を動かすことで，全探索する対象を減らしている．

Root を固定して DFS で求まる． \(i,j\) のうち，root から見て深い方を \(i\) とすると， \begin{align} &\order{ \set{ (i,j) \in V(G)^{2} }{ i \neq j \text{かつ} e \in E(mini\text{-}path_{i,j})} } \\ =& \order{V(G_{i})} \cdot( \order{ V(G) } - \order{V(G_{i})}). \end{align} ここで，\(V(G_{i})\) は \(i\) を root とする \(G\) における subtree.

使っている記号，マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"

int main() {

  ll n;

  cin >> n;

  vvll to(n);

  rep(i,n-1){

    ll a,b; cin >> a >> b;

    --a, --b;

    to[a].push_back(b);

    to[b].push_back(a);

  }

  vll sz(n);

  ll ans = nCr(n,3LL) + nCr(n,2LL);

  auto dfs = [&](auto dfs, ll cu, ll pa = -1) -> void {

    sz[cu] = 1;

    for(auto ne: to[cu]) if(ne != pa){

      dfs(dfs, ne, cu);

      sz[cu] += sz[ne];

    }

    if(pa != -1){

      ans -= sz[cu] * (n - sz[cu]);

    }

  };

  dfs(dfs, 0);

 

  cout << ans << endl;

  return 0;

}