SRM630 div1 easy Egalitarianism3

概要

n頂点の重み付きの木が与えられる。集合の任意の2頂点の距離が等しくなるような頂点集合のうち、集合の要素数の最大を求めろ。

解法

ある頂点を根とした木で根からの距離を考える。根からの距離が等しい2頂点a, bでaとbのLCAが根であるような頂点であれば頂点集合に入れることができる。
根rの子を r[0], r[1], …, r[k] とする。r[i]を根とする部分木で根からの距離dの頂点が存在するならば頂点間の距離が2dであるような頂点集合S(2d)の要素数を1増やせる。r[i]の部分木の頂点の根から距離の集合Tをdfsで求めたあと、S(2d)に+1する。これを全てのr[i]について計算することで、根rから距離が等しくなるような頂点集合の取りうる要素数を求められる。

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
typedef vector<int> VI;
typedef vector<VI> VVI;
typedef vector<ll> VL;
typedef vector<VL> VVL;
typedef pair<int, int> PII;

#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(x) x.begin(), x.end()
#define IN(a, b, x) (a<=x&&x<b)
#define MP make_pair
#define PB push_back
const int INF = (1LL<<30);
const ll LLINF = (1LL<<60);
const double PI = 3.14159265359;
const double EPS = 1e-12;
const int MOD = 1000000007;
//#define int ll

template <typename T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template <typename T> T &chmax(T &a, const T &b) { return a = max(a, b); }

int dx[] = {0, 1, 0, -1}, dy[] = {1, 0, -1, 0};

ll d[55][55];
vector<PII> g[55];
int v[1000010];
set<int> st;

void dfs(int v, int p, int r) {
  for(PII i: g[v]) {
    if(i.first == p) continue;
    st.insert(d[r][i.first]);
    dfs(i.first, v, r);
  }
}

class Egalitarianism3 {
   public:
   int maxCities(int n, vector <int> a, vector <int> b, vector <int> len)
  {
    if(n == 1) return 1;
    REP(i, n) g[i].clear();
    REP(i, n) REP(j, n) {
      if(i == j) d[i][j] = 0; else d[i][j] = INF;
    }
    REP(i, n-1) {
      g[a[i]-1].PB({b[i]-1, len[i]});
      g[b[i]-1].PB({a[i]-1, len[i]});
      d[a[i]-1][b[i]-1] = d[b[i]-1][a[i]-1] = len[i];
    }
    REP(k, n) REP(i, n) REP(j, n) chmin(d[i][j], d[i][k] + d[k][j]);

    int ans = 2;
    REP(i, n) {
      // 頂点iを根としたとき
      memset(v, 0, sizeof(v));
      int ret = 0;
      for(PII j: g[i]) {
        st.clear();
        st.insert(d[i][j.first]);
        dfs(j.first, i, i);
        for(int k: st) {
          if(k >= INF) continue;
          v[k]++;
          chmax(ret, v[k]);
        }
      }
      chmax(ans, ret);
    }

    return ans;
  }
};