CodeChef Snackdown 2015 : Online Elimination Round - Chef and Trees

問題

ツリーが与えられる。1本辺を足して閉路を作るとき、直径を小さくすることが可能か判定せよ。

解法

ツリーの直径を求める。

ツリーの直径が偶数の時、中心となる頂点は1つだけである。中心に隣接する頂点のうち、直径を作るパスに含まれるものが3つ以上あるとき、直径を小さくすることができない。

ツリーの直径が奇数の時、中心となる頂点は2つある。それぞれの中心に隣接する中心を除く頂点のうち、直径を作るパスに含まれる頂点を考える。そのような頂点が、各中心にそれぞれ2つずつ以上隣接している時、直径を小さくすることができない。（この辺は Editorial の図が詳しい）

コード

#include <bits/stdc++.h>
using namespace std;

const int INF = 1e9;
void dfs(int pos, const vector<vector<int>> &g, vector<int> &dist) {
  for (int v : g[pos]) {
    if (dist[v] < INF) continue;
    dist[v] = dist[pos] + 1;
    dfs(v, g, dist);
  }
}

// 木の{中心, （あるなら）中心2, 直径}を返す
tuple<int, int, int> pickCenter(const vector<vector<int>> &g) {
  int N = g.size();
  vector<int> d(N, INF);
  d[0] = 0;
  dfs(0, g, d);
  int farest = 0;  // 頂点0から最も遠い頂点
  for (int i = 0; i < N; ++i)
    if (d[farest] < d[i]) farest = i;

  vector<int> dist(N, INF);  // farestからの距離
  fill(dist.begin(), dist.end(), INF);
  dist[farest] = 0;
  dfs(farest, g, dist);

  int farest2 = farest;  // farest から最も遠い頂点
  for (int i = 0; i < N; ++i)
    if (dist[farest2] < dist[i]) farest2 = i;

  vector<int> dist2(N, INF);  // farest2からの距離
  dist2[farest2] = 0;
  dfs(farest2, g, dist2);
  int diameter = dist[farest2];  //木の直径
  int radius = diameter / 2;
  if (diameter % 2 == 0) {
    for (int i = 0; i < N; ++i)
      if (dist[i] == radius && dist2[i] == radius)
        return make_tuple(i, -1, diameter);
  } else {
    int center1, center2;
    for (int i = 0; i < N; ++i) {
      if (dist[i] == radius && dist2[i] == radius + 1) center1 = i;
      if (dist[i] == radius + 1 && dist2[i] == radius) center2 = i;
    }
    return make_tuple(center1, center2, diameter);
  }

  assert(false);
  return make_tuple(-1, -1, -1);
}

// 直径を作るパスに含まれる頂点以外を削除するDFS
bool remove_dfs(int pos, int radius, const vector<vector<int>> &g,
                vector<int> &dist, vector<bool> &removed) {
  if (g[pos].size() == 1 && dist[pos] != 0) {
    if (dist[pos] == radius) return true;
    assert(dist[pos] < radius);
    removed[pos] = true;
    return false;
  }

  bool will_remain = false;
  for (int v : g[pos]) {
    if (dist[v] < INF) continue;
    dist[v] = dist[pos] + 1;
    if (remove_dfs(v, radius, g, dist, removed)) will_remain = true;
  }
  if (!will_remain) removed[pos] = true;
  return will_remain;
}

bool solve(vector<vector<int>> &g) {
  int N = g.size();
  int center1, center2, diameter;
  tie(center1, center2, diameter) = pickCenter(g);
  cerr << center1 << " " << center2 << " " << diameter << endl;
  if (center2 == -1) {
    // 木の中心が1つだけの時（直径が偶数の時）
    assert(diameter % 2 == 0);

    // 直径に含まれない頂点を削除する
    vector<int> dist(N, INF);
    dist[center1] = 0;
    vector<bool> removed(N, false);
    remove_dfs(center1, diameter / 2, g, dist, removed);

    int neighbors = 0;
    for (int v : g[center1])
      if (!removed[v]) neighbors++;
    assert(neighbors >= 2);
    return neighbors == 2;
  } else {
    //ツリーの中心が2つある時（直径が奇数の時）

    // 中心の間を結ぶ辺を削除して扱いやすくする
    for (auto it = g[center1].begin(); it != g[center1].end(); it++)
      if (*it == center2) {
        g[center1].erase(it);
        break;
      }
    for (auto it = g[center2].begin(); it != g[center2].end(); it++)
      if (*it == center1) {
        g[center2].erase(it);
        break;
      }

    // 直径に含まれない頂点を削除する
    vector<int> dist(N, INF);
    dist[center1] = 0;
    dist[center2] = 0;
    vector<bool> removed(N, false);
    remove_dfs(center1, diameter / 2, g, dist, removed);
    remove_dfs(center2, diameter / 2, g, dist, removed);

    int neighbors1 = 0;
    for (int v : g[center1])
      if (!removed[v]) neighbors1++;
    int neighbors2 = 0;
    for (int v : g[center2])
      if (!removed[v]) neighbors2++;
    return neighbors1 == 1 || neighbors2 == 1;
  }
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int T;
  cin >> T;
  while (T--) {
    int N;
    cin >> N;
    vector<vector<int>> g(N);
    for (int i = 0; i < N - 1; ++i) {
      int a, b;
      cin >> a >> b;
      a--, b--;
      g[a].push_back(b);
      g[b].push_back(a);
    }
    if (N == 2) {
      cout << "NO" << '\n';
      continue;
    }
    if (solve(g)) {
      cout << "YES" << '\n';
    } else {
      cout << "NO" << '\n';
    }
  }
}