Shortest Path [第八回 アルゴリズム実技検定 H]

https://atcoder.jp/contests/past202109-open/tasks/past202109_h

前提知識

• LCA

解説

https://atcoder.jp/contests/past202109-open/submissions/26687334

今回の問題は制約が割と緩いので、基本的な典型アルゴリズムを理解していると解ける。

愚直に考えてみる

条件を満たす都市の組(i,j)が存在するかを判定することが要求されているが、
都市の組については全探索できる制約となっている。
なので、愚直に考えるとすると、都市の組(i,j)を全探索して、各組について距離を求めてXのものがあればYesと答えるようにする。

だが、都市の組(i,j)の全列挙はすでにO(N²)となるので、制約から見ると計算時間はあまり余裕がない。
残る問題は木上の2頂点間の距離を高速に求めるという部分であるが、これは典型アルゴリズムが存在する。

木上の2点間の距離

これを求める方法の1つとしてLCAと累積和を利用する手法がある。
典型だからググって読んでみてと書こうとしたがさっと探して見当たらなかったので、少し詳しく書いておく。
LCAと累積和についてはそれぞれ良質な記事が存在しているので書かない。
もし知らない場合は、検索して学習しよう。
自分はLCAはHL分解を利用したものを使っているが、ややこしいので普通にダブリングで実装したものを使うといい。

さて、木上の2点(a,b)について距離を求める。

1. 木の根をどこでもいいので設定
2. LCAの構築
3. すべての頂点について根からの距離を累積和っぽくDFSで求めておく。配列totとする。

ここまでが事前準備で、利用するときは、

1. 頂点(a,b)からLCAを求める。頂点lcとする
2. 頂点(a,b)間の距離はtot[a] + tot[b] - 2 * tot[lc]となる

細かく見ていこう

手順1

根っこはどこでもいいが、ランダムにしておくと、いじわるなケースを意図せず回避できるかもしれない。
本質的な改善ではないから、普通に0にしとけばいい。

手順2

LCAの構築は、どこかを見て、調べてほしい。
自分はHL分解でやっているが、やや難しいので、ダブリングが楽。
どちらにしろ、ライブラリ化しておくといい。

手順3

手順1で決めた根から累積和で配列totを計算する。

tot[cu] := 根から頂点cuまでの経路の重みの和

累積和と書いているのは、とある頂点cuの子供toへは

tot[to] = tot[cu] + (cuからtoへの重み)

のように計算ができるので、累積和っぽく高速に求めることができるためである。

手順4

これはLCAがあればよい。
頂点aと頂点bのLCAを頂点lcとしておこう。

手順5

すると、頂点aから頂点bへの距離は、tot[a] + tot[b] - 2 * tot[lc]で求められる。
もう少しかみ砕くと、

「頂点a -> 頂点b」は「頂点a -> 頂点lc -> 頂点b」と経路を分解できる。
「頂点a -> 頂点lc」はtot[a] - tot[lc]と計算でき、「頂点lc -> 頂点b」はtpt[b] - tot[lc]と計算できる。
よって、足し合わせるとtot[a] + tot[b] - 2 * tot[lc]となる。

int N, X;
vector<pair<int, int>> E[3010];
//---------------------------------------------------------------------------------------------------
int tot[3010];
void dfs(int cu, int pa = -1) {
    fore(p, E[cu]) if (p.first != pa) {
        tot[p.first] = tot[cu] + p.second;
        dfs(p.first, cu);
    }
}
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> X;
    HLDecomposition hld;
    hld.init(N);
    rep(i, 0, N - 1) {
        int a, b, c; cin >> a >> b >> c;
        a--; b--;
        hld.add(a, b);
        E[a].push_back({ b, c });
        E[b].push_back({ a, c });
    }
    
    hld.build(0);
    dfs(0);

    rep(a, 0, N) rep(b, a + 1, N) {
        int lc = hld.lca(a, b);
        int len = tot[a] + tot[b] - 2 * tot[lc];
        if (len == X) {
            cout << "Yes" << endl;
            return;
        }
    }

    cout << "No" << endl;
}