2024-09-03

AtCoder ABC 369 G - As far as possible (3D, 黄色, 600 点)

Euler Tour して、遅延評価セグ木（区間加算 + 区間 min）に乗せて......という解法を考えていたところに、あまりにもシンプルな想定解法を目にして、感動した！

問題へのリンク

問題概要

頂点数 $N$ の重み付き木が与えられる。 $K = 1, 2, \dots, N$ に対して、次の問に答えよ。

$N$ 個の頂点のうち $K$ 個 $v_{1}, v_{2}, \dots, v_{K}$ を指定したとき、頂点 1 を出発して、頂点 1 に戻ってくるウォークのうち、青木君の指定した $K$ 個の頂点をすべて通るものの最小値を $f(v_{1}, v_{2}, \dots, v_{K})$ とする。

$f(v_{1}, v_{2}, \dots, v_{K})$ の最大値を求めよ。

制約

$2 \le N \le 2 \times 10^{5}$

考えたこと

0-indexed で考える。そして、頂点 0 を根とする根付き木を考える。

$K = 1$ のときは、根から最も遠い頂点 $v_{1}$ （根からの距離を $d_{1}$ とする）を指定するのが最適で、答えは $2d_{1}$ となる。

さて、ここで重要な考察として、 $K = 2, 3, \dots$ のときは、頂点 $v_{1}$ を使うものとしてよいことがわかる。ちゃんとした証明を書くのは面倒だが、概略は次の通り。

$K \ge 2$ で、頂点 $v_{1}$ を使用しない最適解があったとする。

その解におけるウォーク上の頂点のうち、頂点 $v_{1}$ からみたときの最近祖先を $l$ とする。そして、 $K$ 個の頂点のうち、 $l$ の祖先である頂点 $v$ を 1 つ選ぶ。

このとき、 $v$ と $v_{1}$ を入れ替えても、解は悪化しない。

よって、 $K = 2$ のときは、次のように考えればよい。

頂点 0 と頂点 $v_{1}$ を結ぶパス上の頂点を 1 点に縮約し、その頂点を新たな根とする
そうしてできた新しい根付き木において、根から最も遠い頂点を $v_{2}$ とする
このとき、頂点 $v_{1}, v_{2}$ を選ぶのが最適である。

$K = 3$ の場合も上と同様の議論によって、頂点 $v_{1}, v_{2}$ を含むウォークのみ考えればよい。つまり、頂点 $0, v_{1}, v_{2}$ を 1 点に縮約して新たな根としたときの、根から最も遠い頂点を $v_{3}$ とすれば、頂点 $v_{1}, v_{2}, v_{3}$ を選ぶのが最適である。

以降も同様である。

木 DP へ

以上の解法を実装するためには、根から各頂点への距離を Euler Tour 状に頂点を配列してできるセグメント木で実装できると考えていた。

が、TL で目にした想定解法がシンプルだったので、そっちを実装した。

「頂点を選ぶたびに縮約グラフ上で根から最も遠い頂点を求める」とするのではなく、「もとのグラフの各 leaf $v$ について、その leaf が根からの距離が最長となる瞬間の縮約グラフにおいける、leaf と根との距離 $d_{v}$ 」を、もとの木上の DFS で一発で求めてしまうことができるのだ。それができるのであれば、もとのグラフ上での leaf を、上記 $d_{v}$ の値が大きい順にソートして、順に足していけばよい。

たとえば、下図の根付き木について、各 leaf $v$ について、しかるべき時点における縮約グラフの根との距離 $d_{v}$ を（その時点での根となる頂点をカッコ内に）書いている。

さて、この $d_{v}$ を求めるために、次の性質を利用した木 DP ができる。

もとのグラフ上での各頂点 $u$ について、 $u$ を根とする部分木に対して、 $u$ から最も遠い頂点 $L_{u}$ と、その頂点と $u$ との距離 $D_{u}$ を求めておく。

このとき、各頂点 $v$ に対して、各小頂点 $c$ に対して次のように判断できる。

小頂点 $c$ が、 $D_{c}$ の値が最大でないならば、 $d_{L_{c}}$ は頂点 $v$ である
小頂点 $c$ が、 $D_{c}$ の値が最大であるならば、 $d_{L_{c}}$ は頂点 $v$ ではなく、その先祖のいずれかである

詳細は、次のコードを参照。全体の計算量は $O(N \log N)$ となる。

コード

#include <bits/stdc++.h>
using namespace std;
using pll = pair<long long, int>;  // (distance, node)

struct Edge {
    int to;
    long long w;
    Edge(int t = 0, long long w = 0) : to(t), w(w) {}
};
using Graph = vector<vector<Edge>>;

int main() {
    long long N, u, v, w;
    cin >> N;
    Graph G(N);
    for (int i = 0; i < N-1; i++) {
        cin >> u >> v >> w, --u, --v;
        G[u].push_back(Edge(v, w));
        G[v].push_back(Edge(u, w));
    }

    vector<pll> leafs;
    auto dfs = [&](auto dfs, int v, int p) -> pll {
        pll res(0, v);
        for (auto e : G[v]) {
            if (e.to == p) continue;
            auto tmp = dfs(dfs, e.to, v);
            tmp.first += e.w;
            if (tmp.first > res.first) {
                leafs.push_back(res);
                res = tmp;
            } else {
                leafs.push_back(tmp);
            }
        }
        return res;
    };
    auto res = dfs(dfs, 0, -1);
    leafs.push_back(res);

    // Greedy に葉を選んでいく
    sort(leafs.begin(), leafs.end(), greater<pll>());
    long long sum = 0;
    for (int k = 0; k < N; k++) {
        if (k < leafs.size()) sum += leafs[k].first * 2;
        cout << sum << endl;
    }
}