1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94
| int n = edges.size() + 1; vector neibor(n, vector<pair<int, int>>()); for (auto &e : edges) { int i = e[0], j = e[1], w = e[2]; neibor[i].emplace_back(j, w); neibor[j].emplace_back(i, w); }
const int MX = bit_width(n - 1); vector fa(n, vector(MX + 1, 0)); vector cost(n, vector(MX + 1, 0)); vector depth(n, 0);
[&](this auto &&dfs, int u, int father) -> void { for (auto &[v, w] : neibor[u]) { if (v == father) { continue; } depth[v] = depth[u] + 1; fa[v][0] = u; cost[v][0] = w; for (int k = 1; k <= MX; k++) { fa[v][k] = fa[fa[v][k - 1]][k - 1]; cost[v][k] = cost[fa[v][k - 1]][k - 1] + cost[v][k - 1]; } dfs(v, u); } }(1, 0);
auto lca = [&](int u, int v) -> int { if (depth[u] > depth[v]) { swap(u, v); }; int dist = depth[v] - depth[u]; for (int k = 0; dist; ++k, dist >>= 1) { if (dist & 1) { v = fa[v][k]; } } if (u == v) { return u; } for (int k = MX; ~k; k--) { if (fa[u][k] != fa[v][k]) { u = fa[u][k]; v = fa[v][k]; } } return fa[u][0]; };
auto cal_cost = [&](int u, int v) -> int { if (depth[u] > depth[v]) { swap(u, v); }; int dist = depth[v] - depth[u], ans = 0; for (int k = 0; dist; ++k, dist >>= 1) { if (dist & 1) { ans += cost[v][k]; v = fa[v][k]; } } if (u == v) { return ans; } for (int k = MX; ~k; k--) { if (fa[u][k] != fa[v][k]) { ans += cost[u][k] + cost[v][k]; u = fa[u][k]; v = fa[v][k]; } } ans += cost[u][0] + cost[v][0]; return ans; };
|