Official

G - Distance Queries on a Tree Editorial by MMNMM


この問題は、Euler Tour と呼ばれる技法を使って解くことができます。(Euler Tour の詳細についてはここでは解説しません。一言で述べると、根から深さ優先探索を行う際に訪れる辺を向きを含めて順に並べたもの(あるいは頂点を訪れる順に並べたもの)です。)

与えられた木の適当な頂点 \(r\) を根として根付き木とします。

頂点 \(u,v\) の距離を \(d(u,v)\) と書くと、根 \(r\) および \(u,v\) の最小共通祖先 \(\operatorname{lca}(u,v)\) について次の式が成り立ちます。

\[\begin{aligned}d(u,v)&=d(u,\operatorname{lca}(u,v))+d(\operatorname{lca}(u,v),v)\\&=(d(r,u)-d(r,\operatorname{lca}(u,v)))+(d(r,v)-d(r,\operatorname{lca}(u,v)))\\&=d(r,u)+d(r,v)-2d(r,\operatorname{lca}(u,v))\end{aligned}\]

よって、この問題は次の \(2\) 種類のクエリを高速に処理することで解くことができます。

  • 頂点 \(u\) に対して、\(d(r,u)\) を求める。
  • 頂点 \(u,v\) に対して、\(\operatorname{lca}(u,v)\) を求める。

前者は重みの更新があるので同じ \(u\) に対する結果が変わる場合があり、後者は木の形が変わらないので同じ \((u,v)\) に対する結果は変わりません。

Euler Tour を用いることで、それぞれ適当な列について次の種類のクエリを高速に処理すればよいことになります。前者と後者で必ずしも同じ列を管理するわけではないことに注意してください。

  • 重みを更新する・頂点 \(u\) に対して \(d(r,u)\) を求める。
    • 更新を行う辺 \(e\) に対して Euler Tour の \(+e,-e\) に対応する場所の重み \(+w,-w\) を変更する
    • 頂点 \(u\) に対応する先頭からの区間の重みの合計を得る
  • 頂点 \(u,v\) に対して、\(\operatorname{lca}(u,v)\) を求める。
    • 頂点 \(u,v\) に対応する区間における最小値の index を得る

それぞれ Binary Indexed Tree やセグメント木、Sparse Table などを用いることで高速に処理することができます。

たとえば前者を Binary Indexed Tree 、後者をセグメント木で管理すると、空間計算量を \(O(N)\) 、時間計算量を \(O(N+Q\log N)\) とすることができます。

実装例は以下のようになります。

#include <limits>
#include <iostream>
#include <vector>
#include <tuple>
#include <atcoder/segtree>
#include <atcoder/fenwicktree>

// セグメント木の準備
namespace segment_tree_setup {
    using value_type = std::pair<unsigned long, unsigned long>;

    value_type op_min(value_type a, value_type b) { return std::min(a, b); }

    value_type e_inf() {
        return {std::numeric_limits<value_type::first_type>::max(), std::numeric_limits<value_type::second_type>::max()};
    }
}

int main() {
    using namespace std;

    // オイラーツアーを求める
    // depth[i] := (i 番目に訪れる頂点の深さ, i) (1 ≤ i ≤ 2N−1)
    // dist[i] := i 番目に訪れる辺の(向き付けられた)重み (1 ≤ i ≤ 2N−2)
    //   ∑_{i ≤ k} dist[i] = k 番目に訪れる頂点と根との距離
    // edge_to_index[i] := i 番目の辺を (いつおりたか, いつのぼったか) (1 ≤ i ≤ N − 1)
    // vertex_to_index[i] := i 番目の頂点をはじめにいつ訪れたか (1 ≤ i ≤ N)
    const auto[depth, dist, edge_to_index, vertex_to_index]{[] {
        unsigned N;
        cin >> N;
        using edge_type = tuple<unsigned long, unsigned long, unsigned long>;
        vector<vector<edge_type>> edges(N);
        vector<edge_type> edge_index(N - 1);
        {
            unsigned long i{};
            for (auto&& [u, v, c] : edge_index) {
                cin >> u >> v >> c;
                edges[--u].emplace_back(--v, c, i);
                edges[v].emplace_back(u, c, i);
                ++i;
            }
        }

        // オイラーツアーの列 (i 番目に使う辺, i 番目に訪れる頂点)
        vector<pair<unsigned long, unsigned long>> euler_tour;
        euler_tour.reserve(2 * N - 2);
        [dfs{[&edges, &euler_tour](auto &&f, unsigned now, unsigned prev) -> void {
            for (const auto&[next, cost, index] : edges[now])
                if (next != prev) {
                    euler_tour.emplace_back(index, now);
                    f(f, next, now);
                    euler_tour.emplace_back(index, next);
                }
        }}] { dfs(dfs, 0, 0); } ();

        vector<unsigned long> ret_dist, ret_vertex(N, 3 * N);
        ret_vertex[0] = 0;
        vector<pair<unsigned long, unsigned long>> ret_depth{{0, 0}}, ret_edge(N - 1, {3 * N, 3 * N});

        for(unsigned long i{}; i < 2 * N - 2; ++i){
            const auto [edge, vertex]{euler_tour[i]};
            const auto [u, v, c]{edge_index[edge]};
            const auto next_vertex{u ^ v ^ vertex};
            ret_vertex[next_vertex] = min(ret_vertex[next_vertex], i + 1);
            if(ret_edge[edge].first > 2 * N){
                ret_edge[edge].first = i;
                ret_dist.emplace_back(c);
                ret_depth.emplace_back(ret_depth.back().first + 1, ret_depth.back().second + 1);
            }else{
                ret_edge[edge].second = i;
                ret_dist.emplace_back(-c);
                ret_depth.emplace_back(ret_depth.back().first - 1, ret_depth.back().second + 1);
            }
        }

        return make_tuple(ret_depth, ret_dist, ret_edge, ret_vertex);
    }()};

    // 区間最小値を求めるセグメント木
    atcoder::segtree<segment_tree_setup::value_type, segment_tree_setup::op_min, segment_tree_setup::e_inf> lca_depth(depth);

    // 先頭からの和を求める BIT
    atcoder::fenwick_tree<unsigned long> dist_from_root(size(dist));
    for (unsigned i{}; i < size(dist); ++i)
        dist_from_root.add(i, dist[i]);

    unsigned Q;
    cin >> Q;
    for (unsigned i{}; i < Q; ++i) {
        unsigned long t, x, y;
        cin >> t >> x >> y;
        if (t == 1) {
            const auto&[first, last]{edge_to_index[--x]};
            // いまの重みとの差分を求める
            const auto before_cost{dist_from_root.sum(first, first + 1)};
            dist_from_root.add(first, y - before_cost);
            dist_from_root.add(last, -y + before_cost);
        } else {
            const auto u{vertex_to_index[--x]}, v{vertex_to_index[--y]};
            // lca は u, v を含む区間の中の深さ最小の頂点
            const auto lca{lca_depth.prod(min(u, v), max(u, v) + 1).second};
            // dist(u, v) = dist(r, u) + dist(r, v) - 2 dist(r, lca(u, v))
            cout << dist_from_root.sum(0, u) + dist_from_root.sum(0, v) - 2 * dist_from_root.sum(0, lca) << endl;
        }
    }

    return 0;
}

posted:
last update: