题目描述

LCA+dfs+bfs+前缀和

样例

#include <iostream>
#include<algorithm>
#include<cmath>
#include<vector>
#include<queue>
#include<cstring>
using namespace std;
typedef long long ll;
const int N = 5e5 + 10;
struct edge {
    int v, w;
};
int n, m, a, b;
vector<edge> e[N];
int dep[N], fa[N][20];
int k[N];
ll dis[N];
ll ans[N];
bool vis[N];
void dfs(int u, int father) {
    dep[u] = dep[father] + 1;
    fa[u][0] = father;
    for (int i = 1; i <= 19; i++)
    {
        fa[u][i] = fa[fa[u][i - 1]][i - 1];
    }
    for (auto& edge : e[u])
    {
        int v = edge.v;
        if (v != father) {
            dfs(v, u);
        }
    }
}
int lca(int u, int v)
{
    if (dep[u] < dep[v]) {
        swap(u, v);
    }
    //先跳到同一层
    for (int i = 19; i >= 0; i--)
    {
        if (dep[fa[u][i]] >= dep[v]) {
            u = fa[u][i];
        }
    }
    if (u == v) {
        return v;
    }
    //再跳到LCA的下一层
    for (int i = 19; i >= 0; i--)
    {
        if (fa[u][i] != fa[v][i]) {
            u = fa[u][i], v = fa[v][i];
        }
    }
    return fa[u][0];
}
void bfs()
{
    memset(vis, false, sizeof vis);
    queue<int> q;
    q.push(1);
    vis[1] = true;
    while (!q.empty())
    {
        int s1 = q.front();
        q.pop();
        for (auto& edge : e[s1])
        {
            int vv = edge.v;
            int ww = edge.w;
            if (!vis[vv]) {
                dis[vv] = dis[s1] + ww;
                vis[vv] = true;
                q.push(vv);
            }
        }
    }
}
ll redis(int x, int y)
{
    int gether1 = lca(x, y);
    return dis[x] + dis[y] - 2 * dis[gether1];
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    cin >> n >> m;
    for (int i = 1; i < n; i++) {
        int x, y, w;
        cin >> x >> y >> w;
        e[x].push_back({ y,w });
        e[y].push_back({ x,w });
    }
    dfs(1, 0);
    for (int i = 0; i < m; i++) {
        cin >> k[i];
    }
    bfs();
    for (int i = 0; i < m; i++)
    {
        if (!i) {
            ans[i] = 0;
        }
        else {
            ans[i] = ans[i - 1] + redis(k[i], k[i - 1]);
        }
    }
    for (int i = 0; i < m; i++)
    {
        if (!i) {
            cout << ans[m - 1] - ans[1] << " ";
        }
        else if (i == (m - 1)) {
            cout << ans[m - 2] << " ";
        }
        else {
            cout << ans[m - 1] - (ans[i + 1] - ans[i - 1]) + redis(k[i - 1], k[i + 1]) << " ";
        }
    }
    return 0;
}