UOJ-493 CSP-S 2019 树的重心

Description

给定一颗 \(n\) 个顶点的树 \(\text T\)，共 \(n-1\) 次断边操作，每次将树分为两部分 \(\text T_1, \text T_2\)，求：

\[\sum_{(u, v)\in\text E}\left( \sum_{x \text{为} \text T_1 \text{的重心} } x + \sum_{y \text{为} \text T_2 \text{的重心} } y \right) \]

注意，一棵树的重心可能有两个。

Hint

\(1\le n\le 3\times 10^5\)

Solution

不会 \(O(n\log n)\) 或 \(O(n\log^2 n)\) 的做法，这里介绍 \(O(n)\) 的做法 头都写掉了。首先考虑重心的一些性质：

• 一棵树的重心必定存在于 根所在的重链 上。
  • 略证：我们称 重儿子 \(wson(x)\) 为当前根 \(x\) 的子结点中子树大小（\(size\)）最大的那一个。如果 \(size(wson(x)) \le \tfrac 1 2\)，那么根自己就是重心；如果 \(size(wson(x)) > \tfrac 1 2\)，若是选取非重儿子那必然会有一部分的 \(size > \tfrac 1 2\)，因此重心必定在重儿子的子树内。
  • 这有什么用呢？每次找重心可以 只在重链上考虑 而不是整棵树。
• 一颗子树的重心必然为其重儿子子树的重心的一个 祖先。
  • 略证：由上一个性质得，重儿子这个重心、当前子树的重心都在当前子树根所在的重链上，那么加上当前子树根其他部分后重心只会向上偏移，即祖先。
  • 这有什么用呢？这意味着在知道重儿子的重心后，我们不难向上调整出当前新的重心，从而做到 \(O(n)\) 求出 所有子树的重心。
• 一颗树的重心如果存在两个，那么它们 必然相邻。
  • 略证：考虑这样一个性质：以重心为根的树，其根的所有子树大小 不超过 原树的 \(\tfrac 1 2\)。“不超过”包含“小于”和“等于”。如果是小于自然就只有一个重心；如果等于，如果以将树根换成原来的根的 重儿子，又会出现恰好等于一半的情况，新的重儿子其实就是 原来的根。如此重心只会在这两个顶点上移动，那么它们显然是相邻的。
  • 这有什么用呢？实际在操作是就可以优先找 深度相对较大（若在，转为有根树两个相邻重心必然一父一子）的重心，在判断其父亲是不是即可。注意，下面预处理的重心都是深度大的。

以上性质是下面算法的理论基础。

---

我们先令原树的一个重心作为根 \(root\)，然后看看断掉一条边 \((x, y)\) 后会发生什么（\(x\) 为 \(y\) 的父亲）。

原树 \(\text T\) 分裂为两部分：以 \(y\) 为根的子树 \(\text T^\prime\)，以及剩下部分 \(\text T - \text T^\prime\)。

• 对于 \(\text T^\prime\) 部分：由于上面已经提到了 \(O(n)\) 预处理所有子树重心的方法，那么这就没啥问题。
• 对于 \(\text T - \text T^\prime\) 部分：对在原树中删去部分的位置分类讨论：
  • 如果 \(\text T^\prime\) 不在根结点所在重链上，那么原来那个重心会在其所在这条重链上移动。可以预处理出一个 \(\text{ans}_1(s)\) 表示删去大小为 \(s\) 的 非根所在重链部分 后，剩余 \(\text T - \text T^\prime\) 部分的重心是哪个。这个就是跑一遍 \(root\) 所在的重链，复杂度 \(O(n)\)。
  • 如果 \(\text T^\prime\) 在根结点所在重链上，那么 \(root\) 的原来那个重儿子不一定还是重儿子了，对此又分两种情况讨论：
    • 如果重儿子没变，那么原来的重心会 上移，而如果重心本来就是根那必然还是这个根，这就是我们将重心设为 \(root\) 的原因。
    • 如果重儿子变了，首先可以肯定的是，新的重儿子就是原先的 次重儿子，原来的重心会转移到次重子树中。于是也可以预处理一个 \(\text{ans}_2(s)\) 表示删去重链上的大小为 \(s\) 的子树后的重心，具体方法和 \(\text{ans}_1\) 类似，只要在 \(root\) 的位置先走次重儿子即可。复杂度显然也是 \(O(n)\)。

总复杂度显然 \(O(n)\)，但是带一个大常数。

Code

实现细节非常多，必须保证自己思路清晰。

/*
 * Author : _Wallace_
 * Source : https://www.cnblogs.com/-Wallace-/
 * Problem : CSP-S 2019 树的重心
 */
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <functional>
#include <vector>

const int N = 3e5 + 5;
int n, T;
std::vector<int> adj[N];

int siz[N], wson[N], root;
void findRoot(int x, int f) {
  siz[x] = 1, wson[x] = 0;
  for (auto y : adj[x]) if (y != f) {
    findRoot(y, x), siz[x] += siz[y];
    if (siz[wson[x]] < siz[y]) wson[x] = y;
  }
  if (siz[wson[x]] * 2 <= n && (n - siz[x]) * 2 <= n)
    root = x;
}

int dep[N], fa[N], sid[N], swson;
void prework(int x, int f) {
  siz[x] = 1, dep[x] = dep[fa[x] = f] + 1;
  sid[x] = (f == root) ? x : sid[f];
  for (auto y : adj[x]) if (y != f) {
    prework(y, x), siz[x] += siz[y];
    if (siz[wson[x]] < siz[y]) {
      if (x == root) swson = wson[x];
      wson[x] = y;
    } else if (x == root && siz[y] > siz[swson])
      swson = y;
  }
}

int ans1[N], ans2[N], ans3[N], cutSiz;
void calc1(int x, int& cutSiz) {
  if (wson[x]) calc1(wson[x], cutSiz);
  while (cutSiz && cutSiz >= n - siz[x] * 2)
    ans1[cutSiz--] = x;
}
void calc2(int x, int& cutSiz) {
  if (x == root) calc2(swson, cutSiz);
  else if (wson[x]) calc2(wson[x], cutSiz);
  while (cutSiz && cutSiz >= n - siz[x] * 2)
    ans2[cutSiz--] = x;
}
void calc3(int x) {
  if (wson[x]) calc3(wson[x]);
  for (auto y : adj[x]) if (y != fa[x] && y != wson[x]) calc3(y);
  ans3[x] = wson[x] ? ans3[wson[x]] : x;
  while (siz[ans3[x]] * 2 < siz[x]) ans3[x] = fa[ans3[x]];
}

void clear() {
  memset(ans1, 0, sizeof(ans1));
  memset(ans2, 0, sizeof(ans2));
  memset(ans3, 0, sizeof(ans3));

  for (int i = 1; i < N; i++) adj[i].clear();
  memset(dep, 0, sizeof(dep));
  memset(siz, 0, sizeof(siz));
  memset(wson, 0, sizeof(wson));
  memset(fa, 0, sizeof(fa));
  memset(sid, 0, sizeof(sid));
}

void solve() {
  std::function<bool(int, int)> check[3] = {
    [&](int x, int y) {
      return x && siz[wson[x]] * 2 <= siz[y]
          && (siz[y] - siz[x]) * 2 <= siz[y];
    },
    [&](int x, int y) {
      if (x == root) return siz[swson] * 2 <= n - siz[y];
      return x && siz[wson[x]] * 2 <= n - siz[y]
          && (n - siz[x] - siz[y]) * 2 <= n - siz[y];
    },
    [&](int x, int y) {
      if (x == root) return siz[wson[x]] * 2 <= n - siz[y];
      return x && siz[wson[x]] * 2 <= n - siz[y]
          && (n - siz[x] - siz[y]) * 2 <= n - siz[y];
    }
  };

  long long ans = 0;
  for (int i = 1; i <= n; i++) for (auto j : adj[i]) if (i < j) {
    int x = i, y = j;
    if (dep[x] > dep[y]) std::swap(x, y);
    int yc = ans3[y]; ans += yc;

    if (dep[fa[yc]] >= dep[y] && check[0](fa[yc], y))
      ans += fa[yc];

    if (sid[y] == wson[root]) {
      if (siz[wson[root]] - siz[y] >= siz[swson]) {
        ans += root;
      } else {
        int xc = ans2[siz[y]]; ans += xc;
        if (check[1](fa[xc], y)) ans += fa[xc];
      }
    } else {
      int xc = ans1[siz[y]]; ans += xc;
      if (check[2](fa[xc], y)) ans += fa[xc];
    }
  }
  printf("%lld\n", ans);
}

signed main() {
  scanf("%d", &T);
  while (T--) {
    scanf("%d", &n), clear();
    for (int i = 1, u, v; i < n; i++) {
      scanf("%d%d", &u, &v);
      adj[u].emplace_back(v);
      adj[v].emplace_back(u);
    }

    root = 0, findRoot(1, 0);
    memset(wson, 0, sizeof(wson)), swson = 0, prework(root, 0);
    cutSiz = n; calc1(root, cutSiz);
    cutSiz = n, calc2(root, cutSiz);
    calc3(root), solve();
  }
  return 0;
}