QTREE2 - Query on a tree II 题解

目录

• QTREE2 - Query on a tree II
  • 前置知识
  • 定义
  • First. 求 \(dis_{u, v}\)
  • Second. 求 \(u\) 到 \(v\) 路径上的第 \(k\) 个点
  • 时间复杂度
  • Code

QTREE2 - Query on a tree II

\(\mathtt {TAGS}\)：\(\mathtt{LCA}\)，倍增，树上操作

\(\mathtt {APPRAISE}\)：倍增板题。

前置知识

• 最近公共祖先 - OI Wiki
• 最近公共祖先 - 倍增算法

定义

由于根节点是哪一个对本题的求解并没有影响（因为树上两点的简单路径是唯一的），所以本文中默认根节点为 \(1\)。

• \(dis_{u,v}\) 指 \(u\) 到 \(v\) 的唯一简单路径长度。
• \(d_{u}\) 指 \(u\) 到根节点的路径长度。
• \(fa_{u, k}\) 指 \(u\) 的第 \(2 ^ k\) 个祖先。
• \(dep_u\) 指 \(u\) 的深度。

First. 求 \(dis_{u, v}\)

方法 1：

因为 \(u,v\) 不一定在同一颗子树上，直接求 \(dis_{u, v}\) 不好求，那么我们尝试将其转化为我们好求的值 \(d_u\) 和 \(d_v\)，通过他们得到 \(dis_{u,v}\)。

不妨把 \(dis_{u, v}\) 拆解成 \(dis_{u, lca(u,v)}\) 和 \(dis_{\text{lca(u,v)},v}\)，此时由于 \(\text{lca(u,v)}\) 是 \(u\) 的祖先，所以它一定在 \(u\) 到 根节点的路径上，那么 \(dis_u = d_u - d_{\text{lca(u,v)}}\)，同理，\(dis_v = d_v - \text{lca(u,v)}\)。

所以 \(dis_{u,v}\) 就化为了 \(d_u + d_v - d_{\text{lca(u,v)}} \times 2\)。

如下图：

方法 2.
使用倍增在 \(\mathtt{LCA}\) 的过程中求出 \(dis_{u,\text{lca(u,v)}}\) 和 \(dis_{v, \text{lca(u,v)}}\)。

定义 \(ds_{u,k}\) 为 \(u\) 到其第 \(2 ^ k\) 个祖先的简单路径距离（可以预处理）。

在用倍增向上跳时，将两个点跳的距离加在一个 \(tmp\) 上，最后跳至 \(\mathtt{LCA}\) 后，\(tmp\) 即为所求。

这里作者使用的方便一点的方法 1。

方法 2 在此给出代码：

• len 即上文中的 \(ds\)。

void dfs(int u, int f) {
   dep[u] = dep[f] + 1, fa[u][0] = f, len[u][0] = e[u][f];
   for (int i = 1; ; i ++) {
   	fa[u][i] = fa[fa[u][i - 1]][i - 1];
   	len[u][i] = len[fa[u][i - 1]][i - 1] + len[u][i - 1]; // 预处理
   	if (fa[u][i] == 0) {
   		k = max(k, i - 1);
   		break;
   	}
   }
   for (auto v : g[u]){
   	if(v != f){
   		dfs(v, u);
   	}
   }
}
int lca(int u, int v){
   if(dep[u] < dep[v]) swap(u, v);
   int x = u, y = v;
   int lu = 0, lv = 0; // 统计两边跳的距离
   for (int i = k; i >= 0; i --){
   	if(dep[fa[u][i]] >= dep[v])  lu += len[u][i], u = fa[u][i];
   }
   if(u == v) return lu + lv;
   for (int i = k; i >= 0; i --){
   	if(fa[u][i] != fa[v][i]) lu += len[u][i], u = fa[u][i], lv += len[v][i], v = fa[v][i];
   }
   lu += len[u][0], lv += len[v][0];
   return lu + lv;
}

Second. 求 \(u\) 到 \(v\) 路径上的第 \(k\) 个点

这个点的位置有两种情况：

1. 在 \(u\) 到 \(\mathtt{LCA}\) 的路径上。
2. 在 \(\mathtt{LCA}\) 到 \(v\) 的路径上。

对于情况 1:

相当于求 \(u\) 的第 \(k - 1\) 个祖先（因为 \(u\) 为第一个点）

对于情况 2：

转化一下，相当于求 \(v\) 的第 \(\text{路径点数 - k}\) 个祖先。

如下图蓝色路径：

接着直接将 \(k\) 二进制拆分，然后用 \(f_u\) 向上跳至第 \(k\) 祖先即可（具体实现详见代码）。

时间复杂度

• 倍增预处理：\(\text{O}(n \log {n})\)
• \(\mathtt{LCA}\)：\(\text{O}(\log n)\)

总时间复杂度：$\text{O}(t \times (n \log {n}+ q \log n)) $

可以通过此题。

Code

vector< pair <int , int> > G[N];
int f[N][30];
ll dis[N];
int dep[N] , K;
void dfs(int u , int fa) { // 预处理
    dep[u] = dep[fa] + 1;
    f[u][0] = fa;
    for (int i = 1; ; i++) {
        f[u][i] = f[f[u][i - 1]][i - 1];
        if (f[u][i] == 0) {
            K = max(K , i - 1);
            break;
        }
    }
    for (auto e : G[u]) {
        int v = e.first , w = e.second;
        if (v != fa) {
            dis[v] = dis[u] + w;
            dfs(v , u);
        }
    }
}
int lca(int u , int v) { // 倍增求 LCA
    if (dep[u] < dep[v]) swap(u , v);
    for (int i = K; i >= 0; i--) {
        if (dep[f[u][i]] >= dep[v]) u = f[u][i];
    }
    if (u == v) return u;
    for (int i = K; i >= 0; i--) {
        if (f[u][i] != f[v][i]) u = f[u][i] , v = f[v][i];
    }
    return f[u][0];
}
int kth_num(int u , int v , int k) {
    int Lca = lca(u , v);
    int dul = dep[u] - dep[Lca] + 1;
    if (dep[u] - dep[Lca] + 1 >= k) { // 情况 1
        int tmp = u;
        k--; // k - 1 个祖先
        for (int i = log2(k); i >= 0; i--) { // 二进制拆分
            if ((k >> i) & 1) { // 需要跳 2^K 次
                tmp = f[tmp][i];
            }
        }
        return tmp;
    } else {// 情况 2
        k -= dep[u] - dep[Lca] + 1;
        k = dep[v] - dep[Lca] - k;
        int tmp = v;
        for (int i = log2(k); i >= 0; i--) {// 二进制拆分
            if ((k >> i) & 1) {// 需要跳 2^K 次
                tmp = f[tmp][i];
            }
        }
        return tmp;
    }
}

signed main() {
    ios::sync_with_stdio(0);
    int t;
    cin >> t;
    while (t--) {
        cin >> n;
        for (int i = 1; i < n; i++) {
            int u , v , w;
            cin >> u >> v >> w;
            G[u].push_back({ v, w });
            G[v].push_back({ u, w });
        }
        dfs(1 , 0);
        string opt;
        int u , v , k;
        while (cin >> opt) {
            if (opt == "DONE") break;
            if (opt == "DIST") {
                cin >> u >> v;
                int Lca = lca(u , v);
                cout << dis[u] + dis[v] - dis[Lca] * 2 << endl;
            } else {
                cin >> u >> v >> k;
                cout << kth_num(u , v , k) << endl;
            }
        }
        // 多测一定要清空！！！！！！
        // 本人在此 RE + WA 了三次
        for (int i = 1; i <= n; i++) G[i].clear();
        memset(f , 0 , sizeof f);
        memset(dis , 0 , sizeof dis);
        memset(dep , 0 , sizeof dep);
    }
    return 0;
}