【距离-LCA】

题目

代码

#include <bits/stdc++.h>
using namespace std;

const int N = 1e4+10, M = 2e4+10;
int h[N], e[M], ne[M], w[M], idx;
int d[N], p[N][20][2];
int n, m;

// 添加一条从 a 到 b 的边，权重为 c
void add(int a, int b, int c)
{
    e[idx] = b, w[idx] = c, ne[idx] = h[a], h[a] = idx++;
}

// DFS 遍历树，记录每个节点的深度以及父节点
void dfs(int from, int u)
{
    for(int k = h[u]; ~k; k = ne[k])
    {
        int to = e[k];
        if(to == from) continue;
        d[to] = d[u] + 1;
        p[to][0][0] = u;
        p[to][0][1] = w[k];
        dfs(u, to);
    }
}

// 预处理 LCA
void init()
{
    for(int i = 1; i <= 16; i++)
    {
        for(int j = 1; j <= n; j++)
        {
            p[j][i][0] = p[p[j][i-1][0]][i-1][0];
            p[j][i][1] = p[j][i-1][1] + p[p[j][i-1][0]][i-1][1];
        }
    }
}

// 查询两点之间的最短距离
int lca(int a, int b)
{
    int distance = 0;

    // 保证 a 的深度大于等于 b 的深度
    if(a < b) swap(a, b);

    // 将 a 下移到与 b 同一深度
    for(int i = 16; i >= 0; i--)
    {
        if(d[p[a][i][0]] >= d[b])
        {
            distance += p[a][i][1];
            a = p[a][i][0];
        }
    }

    // 如果此时 a 和 b 相等，则直接返回距离
    if(a == b) return distance;

    // 否则，将 a 和 b 上移至最近公共祖先
    for(int i = 16; i >= 0; i--)
    {
        if(p[a][i][0] != p[b][i][0])
        {
            distance = distance + p[a][i][1] + p[b][i][1];
            a = p[a][i][0];
            b = p[b][i][0];
        }
    }

    // 最后加上最后一跳的距离
    distance += p[a][0][1] + p[b][0][1];

    return distance;
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    memset(h, -1, sizeof h);
    cin >> n >> m;

    for(int i = 1; i <= n-1; i++)
    {
        int x, y, k;
        cin >> x >> y >> k;
        add(x, y, k);
        add(y, x, k);
    }

    d[1] = 1;
    dfs(0, 1);
    init();

    for(int i = 1; i <= m; i++)
    {
        int x, y;
        cin >> x >> y;
        cout << lca(x, y) << '\n';
    }

    return 0;
}

我以为是distance计算的问题，于是改成了下面这版，但还是不对。样例两版都过了。

#include <bits/stdc++.h>
using namespace std;

const int N = 1e4+10, M = 2e4+10;
int h[N], e[M], ne[M], w[M], idx;
int d[N], p[N][20][2], dist[N];
int n, m;

// 添加一条从 a 到 b 的边，权重为 c
void add(int a, int b, int c)
{
    e[idx] = b, w[idx] = c, ne[idx] = h[a], h[a] = idx++;
}

// DFS 遍历树，记录每个节点的深度以及父节点
void dfs(int from, int u)
{
    for(int k = h[u]; ~k; k = ne[k])
    {
        int to = e[k];
        if(to == from) continue;
        d[to] = d[u] + 1;
        p[to][0][0] = u;
        p[to][0][1] = w[k];
        dist[to] = dist[u] + w[k];
        dfs(u, to);
    }
}

// 预处理 LCA
void init()
{
    for(int i = 1; i <= 16; i++)
    {
        for(int j = 1; j <= n; j++)
        {
            p[j][i][0] = p[p[j][i-1][0]][i-1][0];
            p[j][i][1] = p[j][i-1][1] + p[p[j][i-1][0]][i-1][1];
        }
    }
}

// 查询两点之间的最短距离
int lca(int a, int b)
{
    // 保证 a 的深度大于等于 b 的深度
    if(a < b) swap(a, b);

    // 将 a 下移到与 b 同一深度
    for(int i = 16; i >= 0; i--)
    {
        if(d[p[a][i][0]] >= d[b])
        {
            a = p[a][i][0];
        }
    }

    // 如果此时 a 和 b 相等，则直接返回距离
    if(a == b) return a;

    // 否则，将 a 和 b 上移至最近公共祖先
    for(int i = 16; i >= 0; i--)
    {
        if(p[a][i][0] != p[b][i][0])
        {
            a = p[a][i][0];
            b = p[b][i][0];
        }
    }

    return p[a][0][0];
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    memset(h, -1, sizeof h);
    cin >> n >> m;

    for(int i = 1; i <= n-1; i++)
    {
        int x, y, k;
        cin >> x >> y >> k;
        add(x, y, k);
        add(y, x, k);
    }

    d[1] = 1;
    dfs(0, 1);
    init();

    for(int i = 1; i <= m; i++)
    {
        int x, y;
        cin >> x >> y;
        int anc = lca(x, y);
        cout << dist[x] + dist[y] - 2 * dist[anc] << '\n';
    }

    return 0;
}