第八回 アルゴリズム実技検定 H問題

PAST008H

LCAの前計算で高速化可能だが， 頂点毎にDFSをやりなおしても \( O(N^2) \) で間に合う．
Tree においては，根からDFSしたときの距離が，根からの最短距離になる．

使っている記号，マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"

struct Edge {

    long long to;

    long long cost;

  Edge(long long to) : to(to){}

  Edge(long long to, long long co) : to(to), cost(co){}

};

using Graph = vector<vector<Edge>>;

/* LCA(G, root): 木 G に対する根を root として Lowest Common Ancestor を求める構造体

    query(u,v): u と v の LCA を求める。計算量 O(logn)

    前処理: O(nlogn)時間, O(nlogn)空間

*/

struct LCA {

    vector<vector<int>> parent;  // parent[k][u]:= u の 2^k 先の親

    vector<int> dist;            // root からの距離

    LCA(const Graph &G, int root = 0) { init(G, root); }

    // 初期化

    void init(const Graph &G, int root = 0) {

        int V = G.size();

        int K = 1;

        while *1;

        dist.assign(V, -1);

        dfs(G, root, -1, 0);

        for (int k = 0; k + 1 < K; k++) {

            for (int v = 0; v < V; v++) {

                if (parent[k][v] < 0) {

                    parent[k + 1][v] = -1;

                } else {

                    parent[k + 1][v] = parent[k][parent[k][v]];

                }

            }

        }

    }

    // 根からの距離と1つ先の頂点を求める

    void dfs(const Graph &G, int v, int p, int d) {

        parent[0][v] = p;

        dist[v] = d;

        for (auto e : G[v]) {

            if (e.to != p) dfs(G, e.to, v, d + 1);

        }

    }

    int query(int u, int v) {

        if (dist[u] < dist[v]) swap(u, v);  // u の方が深いとする

        int K = parent.size();

        // LCA までの距離を同じにする

        for (int k = 0; k < K; k++) {

            if *2;

    g[b].push_back(Edge(a, c));

  }

  LCA lca(g);

  vll dis(n);

  auto dfs = [&](auto f, ll cu, ll pa = -1, ll d = 0) -> void {

    dis[cu] = d;

   

    for(auto e: g[cu]) if(e.to != pa){

      f(f, e.to, cu, d + e.cost);

    }

  };

  dfs(dfs,0);

 

  rep(i,n) rep(j,i){

    ll l = lca.query(i,j);

    if(x == dis[i] + dis[j] - 2*dis[l]){

      cout << "Yes" << endl;

      return 0;

    }

  }

  cout << "No" << endl;

  return 0;

}

*1:1 << K) < V) K++;

        parent.assign(K, vector<int>(V, -1

*2:dist[u] - dist[v]) >> k & 1) {

                u = parent[k][u];

            }

        }

        // 二分探索で LCA を求める

        if (u == v) return u;

        for (int k = K - 1; k >= 0; k--) {

            if (parent[k][u] != parent[k][v]) {

                u = parent[k][u];

                v = parent[k][v];

            }

        }

        return parent[0][u];

    }

};

int main() {

  ll n, x, m;

  cin >> n >> x;

  m = n-1;

  Graph g(n);

  rep(i,m){

    ll a,b,c; cin >> a >> b >> c;

    --a,--b;

    g[a].push_back(Edge(b, c