NJPC2017 E - 限界集落

問題ページ
E: 限界集落 - NJPC2017 | AtCoder

うしさんのブログを見て前半2問で全方位木DPを学んだので解説を見ずに解いた

考えたこと

• 各頂点から最も遠い頂点までの距離を求めて全てDより大きければ到達不可能なので-1
• d[i] = 頂点iを根とする部分木で頂点iに向かうのに辺が反対の本数 をもちたくなる
• とりあえずdfsすればこれは求まる
• ans[i] = 頂点iに向かうのに辺が反対の本数 を最終的に求めたい
• 頂点iのそれぞれの子vについて (i→v)への辺の方向とd[v] からvを根とする部分木からiに向かうのに辺が反対の本数がわかる
• 親pからはpからiの方向に反対となる辺の本数をd_parとして渡せればよさそう
• 子のd[v]と親からのd_parが適切に与えられればans[i]が計算できる
• p→iに遷移するときd_parはans[p]からiの部分木で反対方向の辺の本数を引いた数にするとよさそう

学び

• 全方位木DPの実装は別に簡単だった
• 各頂点を根にしてO(n^2)みたいなDPが生えたら全方位木DPを考えよう

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
#define int ll
using VI = vector<int>;
using VVI = vector<VI>;
using PII = pair<int, int>;

#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 PB push_back

const ll LLINF = (1LL<<60);
const int INF = (1LL<<30);
const int MOD = 1000000007;

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); }
template <typename T> bool IN(T a, T b, T x) { return a<=x&&x<b; }
template<typename T> T ceil(T a, T b) { return a/b + !!(a%b); }
template<class S,class T>
ostream &operator <<(ostream& out,const pair<S,T>& a){
  out<<'('<<a.first<<','<<a.second<<')';
  return out;
}
template<class T>
ostream &operator <<(ostream& out,const vector<T>& a){
  out<<'[';
  REP(i, a.size()) {out<<a[i];if(i!=a.size()-1)out<<',';}
  out<<']';
  return out;
}

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

vector<PII> g[100010];
map<PII, int> mp;
signed main(void)
{
  int n, d;
  cin >> n >> d;
  REP(i, n-1) {
    int a, b, c;
    cin >> a >> b >> c;
    a--, b--;
    g[a].PB({b, c});
    g[b].PB({a, c});
    mp[{a, b}] = 1;
  }

  VI dist(n), edge(n);
  function<int(int,int)> dfs1 = [&](int v, int p) -> int {
    if(p != -1 && g[v].size() == 1) return 0;
    for(PII &i: g[v]) if(i.first != p) {
      edge[v] += (mp.find(make_pair(v, i.first)) != mp.end());
      edge[v] += dfs1(i.first, v);
      chmax(dist[v], dist[i.first] + i.second);
    }
    return edge[v];
  };

  dfs1(0, -1);

  // 頂点iから最も遠い頂点までの距離
  VI len(n);
  function<void(int,int,int)> dfs2 = [&](int v, int d_par, int p) {
    vector<PII> d_child;
    d_child.emplace_back(0, -1);
    for(PII &e : g[v]) {
      if(e.first == p) d_child.emplace_back(d_par + e.second, e.first);
      else d_child.emplace_back(dist[e.first] + e.second, e.first);
    }
    sort(d_child.rbegin(), d_child.rend());
    len[v] = d_child[0].first;
    for(PII &e : g[v]) {
      if(e.first == p) continue;
      dfs2(e.first, d_child[d_child[0].second == e.first].first, v);
    }
  };

  dfs2(0, -1, -1);

  bool flag = false;
  REP(i, n) if(len[i] <= d) flag = true;
  if(!flag) {
    cout << -1 << endl;
    return 0;
  }

  // 頂点iに向かうのに辺を反転しないといけない本数
  VI ans(n), d_child(n);
  function<void(int,int,int)> dfs3 = [&](int v, int d_par, int p) {
    for(PII &e : g[v]) {
      int tmp = (mp.find(make_pair(v, e.first)) != mp.end());
      if(e.first == p) {
        ans[v] += d_par + tmp;
        d_child[e.first] = d_par + tmp;
      } else {
        ans[v] += edge[e.first] + tmp;
        d_child[e.first] = edge[e.first] + tmp;
      }
    }
    for(PII &e : g[v]) {
      if(e.first == p) continue;
      // ans[v] から e.firstの部分木方向のを引いたやつ
      dfs3(e.first, ans[v] - d_child[e.first], v);
    }
  };

  dfs3(0, 0, -1);

  int ret = INF;
  REP(i, n) {
    if(len[i] <= d) chmin(ret, ans[i]);
  }
  cout << ret << endl;

  return 0;
}