Official

G - Distance Queries on a Tree Editorial by en_translator


This problem can be solved with a technique called Euler Tour. (We do not describe the details on the Euler Tour. Simply put, it is a sequence of edges in the order of traverses during a Depth-First Search (DFS) from the root, including the directions (or a sequence of vertices in the order of visits).

Consider the given tree as a rooted tree rooted at an arbitrary vertex \(r\).

Let \(d(u,v)\) be the distance between vertices \(u\) and \(v\); then the root \(r\), vertices \(u\) and \(v\), and their lowest common ancestor \(\operatorname{lca}(u,v)\) satisfies:

\[\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}\]

Therefore, this problem can be solved by processing the following two queries fast:

  • For vertices \(u\), find \(d(r,u)\).
  • For vertices \(u,v\), find \(\operatorname{lca}(u,v)\).

The result for the former may change for the same \(u\) because the weights are updated, while the latter does not change for the same \(u\) and \(v\).

Using the Euler Tree, it is sufficient to process the following queries for any sequences. Note that the former and the latter may not necessarily manage the same sequences.

  • Updating a weight & finding \(d(r,u)\) for a vertex \(u\):
    • For an edge \(e\) to be updated, change the weights \(+w\) and \(-w\) at the positions corresponding to \(+e\) and \(-e\) in the Euler tree.
    • Find the sum of weights in the segment that starts from the beginning and corresponds to vertex \(u\).
  • Finding \(\operatorname{lca}(u,v)\) for vertices \(u\) and \(v\):
    • Find the index of minimum value that corresponds to vertices \(u,v\).

They can be processed fast enough with a binary indexed tree (Fenwick Tree), segment tree, or sparse table.

For example, managing the former with a binary indexed tree and the latter with a segment tree requires a spacial complexity of \(O(N)\) and a time complexity of \(O(N+Q\log N)\).

The following is a sample code.

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

// Setup of the segment tree
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;

    // Find the Euler Tour
    // depth[i] := (depth of vertex of i-th visit, i) (1 ≤ i ≤ 2N−1)
    // dist[i] := (directed) weight of the edge of i-th traverse (1 ≤ i ≤ 2N−2)
    //   ∑_{i ≤ k} dist[i] = distance between root and vertex of k-th visit
    // edge_to_index[i] := when it (climbed, descended) i-th edge? (1 ≤ i ≤ N − 1)
    // vertex_to_index[i] := when it visited i-th  vertex for the first time? (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;
            }
        }

        // Sequence of Euler tour (i-th traversed edge, i-th visited vertex)
        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);
    }()};

    // segment tree for segment minimum
    atcoder::segtree<segment_tree_setup::value_type, segment_tree_setup::op_min, segment_tree_setup::e_inf> lca_depth(depth);

    // Fenwick tree for prefix sum
    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]};
            // Find the difference from the current weight
            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 is vertex with minimum depth in segment containing u and 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: