题解：AT_abc221_f [ABC221F] Diameter set

题意

给定一棵大小为 的树，现在要选取一些点，使得这些点间的距离均为直径长度。求选取方案数。

思路

由直径的性质：直径具有共同的中点，我们可以从中点下手。

若直径长度为偶数，要使每个节点的距离均为直径长度 ，那么任意两点在以中点为根的情况下的 LCA 都应该为这个中点。且选取的点均为深度为 的叶子节点。

那么，设中点为 ，那么 的每个儿子子树中，均只能选择一个节点。我们可以记录以 为根的子树中满足要求的叶子节点数量 ，然后，统计答案时，将 的儿子子树答案 相乘。

若直径长度为奇数，那么我们可以将树分为两部分，一部分是直径的中边的左端点，另一部分是右端点。

但是，我们也可以在每条边中间都加一个节点，这样可以令直径长度均为偶数，且不会影响答案。

注意，最后的答案需要减去 ，因为答案多出了不选和只选 1 个的情况。

代码

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

const int MAXN = 2e5 + 10, MOD = 998244353;

int n, d[MAXN << 1], pre[MAXN << 1], cnt[MAXN << 1], len;
vector<int> g[MAXN << 1];

void add_edge(int u, int v) {
  g[u].push_back(v);
  g[v].push_back(u);
}

void dfs(int u, int fa) {
  pre[u] = fa;
  for (int v : g[u]) {
    if (v == fa) continue;
    d[v] = d[u] + 1;
    dfs(v, u);
  }
}

void dfs(int u, int fa, int d) {
  cnt[u] = d == len;
  for (int v : g[u]) {
    if (v == fa) continue;
    dfs(v, u, d + 1);
    cnt[u] += cnt[v];
  }
}

int main() {
  ios::sync_with_stdio(0), cin.tie(0);
  cin >> n;
  for (int i = 1, u, v; i < n; i++) {
    cin >> u >> v;
    add_edge(u, i + n);
    add_edge(i + n, v);
  }
  dfs(1, 0);
  int k = max_element(d + 1, d + n + 1) - d;
  d[k] = 0;
  dfs(k, 0);
  k = max_element(d + 1, d + n + 1) - d;
  len = *max_element(d + 1, d + n + 1) >> 1;
  for (int i = 1; i <= len; k = pre[k], i++);
  dfs(k, 0, 0);
  ll ans = 1;
  for (int v : g[k]) {
    (ans *= cnt[v] + 1) %= MOD;
  }
  ans = (ans - 1 + MOD) % MOD;
  for (int v : g[k]) {
    ans = (ans - cnt[v] + MOD) % MOD;
  }
  cout << ans;
  return 0;
}