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
| int n = edges.size() + 1; vector neibor(n + 1, vector<pair<int, int>>()); for (auto &e : edges) { int i = e[0] + 1, j = e[1] + 1, w = e[2]; neibor[i].emplace_back(j, w); neibor[j].emplace_back(i, w); }
auto getUpLog = [](int n) -> int { return n ? __lg(n - 1) + 1 : 0; };
const int MX = getUpLog(n - 1); vector fa(n + 1, vector(MX + 1, 0)); vector cost(n + 1, vector(MX + 1, 0)); vector depth(n + 1, 0);
{ auto dfs = [&](this auto &&dfs, int u, int father) -> void { fa[u][0] = father; depth[u] = depth[father] + 1; for (int k = 1; k <= MX; k++) { fa[u][k] = fa[fa[u][k - 1]][k - 1]; cost[u][k] = cost[fa[u][k - 1]][k - 1] + cost[u][k - 1]; } for (auto &[v, w] : neibor[u]) { if (v != father) { cost[v][0] = w; dfs(v, u); } } }; dfs(1, 0); }
auto lca = [&](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 >= 0 && u != v; 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; };
|