Official

F - Expensive Expense Editorial by Nyaan


この問題は次の 2 つの解き方があります。

  1. 直径 を利用した解き方
  2. 全方位木 DP を利用した解き方

の順に触れていきたいといきます。

1. 直径を利用した解き方

以下の説明ではグラフは木であるとします。また、無向グラフの頂点の組 \((s,t)\) に対して \(s\)\(t\) の距離 (最短パスの長さ) を \(d(s,t)\) のように表します。

直径 とはグラフ内で最も離れている \(2\) 頂点の距離のことをいいます。形式的には、グラフの頂点の組 \((s,t)\) に対する \(d(s,t)\) の最大値とも言えます。

最短距離が直径と一致する頂点の組を \((s,t)\) とします。このときグラフ内の全ての頂点に対して以下の性質が成り立ちます。

すべてのグラフ内の頂点 \(u\) に対して、 \(u\) から最も遠い頂点のうちの 1 つは \(s\)\(t\) である。言い換えると、

\[\max_{v \in G} d(u,v) = \max(d(u,s), d(u,t))\]

が成り立つ。

この問題は直径の性質を利用して次のように解くことができます。

与えられたグラフを \(G(V,E)\) のように置きます。( \(V\) は頂点集合、 \(E\) は辺集合です。)
すべての頂点 \(v \in V\) に対して、補助的なノード \(v'\) を新たに置き、 \(v\)\(v'\) の間に重み \(D_v\) の辺を貼った拡張グラフを新たに考えます。

以下の図は サンプル1 の拡張グラフを図示したものです。

sample1.png

求めるものはすべての \(v \in V\) に対して、

\[\max_{u \in V, u \neq v} d(v,u') \ \cdots (\ast)\]

となりますが、これは拡張グラフの直径を計算して \(d(s',t')\) が直径と一致する補助的なノードの組を \(s',t'\) としたときに

\[ (\ast) = \begin{cases} d(v, t') & v = s \\ d(v, s') & v = t \\ \max(d(v,s'),d(v,t')) & \mathrm{otherwise} \end{cases} \]

であることが、直径の証明と同様の方法で証明できます。

直径はダイクストラ法を \(2\) 回適用すれば計算できます。以上よりこの問題を \(\mathrm{O}(N \log N)\) で解くことができました。

  • 追記:グラフは木なので BFS/DFS でも直径を求めることができます。よってこの問題を \(\mathrm{O}(N)\) で解くことができました。

C++ による想定解は次の通りです。

#include <iostream>
#include <queue>
#include <utility>
#include <vector>
using namespace std;

using ll = long long;
#define rep(i, n) for (int i = 0; i < (n); i++)

struct E {
  int u, v, cost;
};
vector<vector<E>> g;

vector<ll> dijkstra(int start = 0) {
  int N = g.size();
  vector<ll> d(N, 1e18);
  using P = pair<ll, int>;
  priority_queue<P, vector<P>, greater<P>> Q;

  Q.emplace(0, start);
  d[start] = 0;

  while (!Q.empty()) {
    auto [dc, c] = Q.top();
    Q.pop();
    if (d[c] != dc) continue;
    for (auto& e : g[c]) {
      ll dist = dc + e.cost;
      if (dist < d[e.v]) {
        d[e.v] = dist;
        Q.emplace(dist, e.v);
      }
    }
  }

  return d;
}

int main() {
  int N;
  cin >> N;
  g.resize(N);
  rep(i, N - 1) {
    int a, b, c;
    cin >> a >> b >> c;
    --a, --b;
    g[a].push_back(E{a, b, c});
    g[b].push_back(E{b, a, c});
  }
  vector<ll> D(N);
  for (auto& x : D) cin >> x;

  // 直径を求める
  int s = 0, t = -1, u = -1;

  {
    auto ds = dijkstra(s);
    // 最遠点は?
    ll cost = -1;
    rep(i, N) {
      if (i == s) continue;
      if (D[i] + ds[i] > cost) {
        cost = D[i] + ds[i];
        t = i;
      }
    }
  }

  {
    auto dt = dijkstra(t);
    // 最遠点は?
    ll cost = -1;
    rep(i, N) {
      if (i == t) continue;
      if (D[i] + dt[i] > cost) {
        cost = D[i] + dt[i];
        u = i;
      }
    }
  }

  // t, u が間の距離が直径に
  auto dt = dijkstra(t);
  auto du = dijkstra(u);
  rep(i, N) {
    ll cost = -1;
    if (i != t) cost = max(cost, dt[i] + D[t]);
    if (i != u) cost = max(cost, du[i] + D[u]);
    cout << cost << "\n";
  }
}


2. 全方位木 DP を利用した解き方

この問題は 全方位木 DP (Rerooting) と呼ばれるアルゴリズムでも解くことができます。

全方位木 DP の解説は ei13333 氏による記事(日本語) が非常に有名で、全方位木 DP がどのようなアルゴリズムかを知らない人はこちらを読むことをお勧めします。

ここではアルゴリズムを簡潔に述べます。まず、 与えられたグラフを頂点 \(1\) を根とする根付き木とみなします。 \(m_i\)

  • \(m_i:=\) \(i\) から \(i\) の子孫 (\(i\) 自身を含まない) へ旅行したときの旅費の最大

と定義すると、これは次に書く DP で計算できます。

# c:現在の頂点 p:親の頂点
def dfs(c, p):
  dist = 0
  for d in (children of c):
    dfs(d, c)
    dist = max(D[d], m[d]) + (cost of path c-d)
  m[c] = dist
  return

次に前計算を利用して Rerooting を行い、

  • \(dp_i ;=\) \(i\) に対する求める答え

を計算します。

# c: 現在の頂点 p: 親の頂点 
# v: c を根としたときに、親の頂点を根とする部分木から得られる寄与
def dfs2(c, p, v):
  mc = [max(m[d], D[d]) + (cost of edge d-c) for d in (chlidren of c)]
  dp[c] = max(v, max(mc))
  for d in (children of c):
    # d を根とする部分木の寄与のみ取り除く
    exc_d = set(children of c) - d
    mc2 = [max(m[d], D[d]) + (cost of edge d-c) for d in exc_d]
    dv = max(v, max(mc2))
    dfs2(d, c, max(dv, D[c]) + (cost of edge d-c))
  return

2回目の DFS は上記の通りに実装すると \(\mathrm{O}(N^2)\) かかってしまいますが、累積和などのデータ構造を利用すれば 計算量を \(\mathrm{O}(N)\) に落とすことができます。よってこの問題を \(\mathrm{O}(N)\) で解くことができました。

解答例では全方位木 DP を抽象化した実装を行いました。競技プログラミングにおける “抽象化” とは、関数や変数を与えるとアルゴリズムが動作するようなライブラリを作成することで、コンテストの際にアルゴリズムの部分をブラックボックス化して実装できるようにすることを言います。

抽象化のモチベーションとしては、ライブラリを使用してコンテスト中の実装を省略することでより素早く AC を取りたい、というものが主だったものとして挙げられます。たとえば全方位木 DP を利用する問題を例に挙げると、ほとんどの場合は

  • DP テーブルに載せる値の型
  • 頂点同士をマージする関数およびその単位元
  • 親に値を返す関数
  • 子に頂点が存在しないときの値 (この問題では単位元と一致)

の 4 つが問題によって異なり、それ以外は問題によらず同じ実装をすれば問題を解くことができます。そこで、上記の 4 つを与えると自動的に全方位木 DP を計算して計算結果を返してくれるようなライブラリを持っていると確かに高速に AC を取ることができます。

競技プログラミングにおける抽象化の例としては AtCoder Library の Segment Tree などが挙げられます。ドキュメント に詳しい使い方が書かれているので、アルゴリズムを抽象化する例として読んでみるとよいでしょう。

とはいえ、 AtCoder では多大な実装を要求する問題はほとんど存在しないため、ライブラリの作成がレートにもたらす寄与は大きくありません。また、抽象化ライブラリの作成はコンテストで問題を解くよりも得てして大変なことが多いです。

以上の理由から競技プログラミングの勉強法としては不向きかもしれませんが、アルゴリズムやプログラミング言語に対する理解を深める方法の 1 つとしてはありだと個人的に思っています。

C++ による抽象化全方位木 DP の実装例、およびこの問題の解答例は次の通りです。

#include <iostream>
#include <vector>
using namespace std;

template <typename Cost>
struct Edge {
  int src, to;
  Cost cost;
  Edge(int s, int t, Cost c = 1) : src(s), to(t), cost(c) {}
  operator int() const { return to; }
};
template <typename Cost>
struct Graph : vector<vector<Edge<Cost>>> {
  Graph(int n) : vector<vector<Edge<Cost>>>(n) {}
  void add_edge(int s, int t, Cost c = 1) { (*this)[s].emplace_back(s, t, c); }
};

template <typename Cost, typename Data, Data (*merge)(Data, Data), Data (*e)(),
          Data (*leaf)(), Data (*apply)(Data, int, int, Cost)>
struct Rerooting : Graph<Cost> {
  vector<Data> dp, memo;

  Rerooting(int n) : Graph<Cost>::Graph(n) {}

  vector<Data> run() {
    memo.resize((*this).size(), e());
    dp.resize((*this).size());
    dfs1(0, -1);
    dfs2(0, -1, e());
    return dp;
  }

  void dfs1(int c, int p) {
    bool upd = false;
    for (Edge<Cost> &d : (*this)[c]) {
      if (d == p) continue;
      dfs1(d, c);
      upd = true;
      memo[c] = merge(memo[c], apply(memo[d], d, c, d.cost));
    }
    if (!upd) memo[c] = leaf();
  }

  void dfs2(int c, int p, const Data &val) {
    vector<Data> ds{val};
    for (Edge<Cost> &d : (*this)[c]) {
      if (d == p) continue;
      ds.push_back(apply(memo[d], d, c, d.cost));
    }
    int n = ds.size(), idx = 1;
    vector<Data> head(n + 1, e()), tail(n + 1, e());
    for (int i = 0; i++ < n;) head[i] = merge(head[i - 1], ds[i - 1]);
    for (int i = n; i-- > 0;) tail[i] = merge(tail[i + 1], ds[i]);
    dp[c] = head[n];
    for (Edge<Cost> &d : (*this)[c]) {
      if (d == p) continue;
      Data sub = merge(head[idx], tail[idx + 1]);
      dfs2(d, c, apply(sub, c, d, d.cost));
      idx++;
    }
  }
};

vector<int> D;

using Data = long long;
using Cost = long long;
Data merge(Data a, Data b) { return max(a, b); }
Data e() { return 0; }
Data leaf() { return 0; }
Data apply(Data a, int c, int, Cost w) { return max<Data>(a, D[c]) + w; }

int main() {
  int N;
  cin >> N;
  Rerooting<Cost, Data, merge, e, leaf, apply> g(N);
  for (int i = 0; i < N - 1; i++) {
    int s, t, c;
    cin >> s >> t >> c;
    g.add_edge(s - 1, t - 1, c);
    g.add_edge(t - 1, s - 1, c);
  }
  D.resize(N);
  for (auto &x : D) cin >> x;
  auto ans = g.run();
  for (auto &x : ans) cout << x << "\n";
}

posted:
last update: