#include <bits/stdc++.h>
//#pragma GCC target("avx2")
//#pragma GCC optimize("O3,unroll-loops")
#define fastIO ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0)
#define endl '\n'
#define all(v) (v).begin(), (v).end()
#define rall(v) (v).rbegin(), (v).rend()
using namespace std;
using ll [[maybe_unused]] = long long;
using ull [[maybe_unused]] = unsigned long long;
using lll [[maybe_unused]] = __int128;
using lld [[maybe_unused]] = long double;
using llld [[maybe_unused]] = __float128;
template<typename T> using graph [[maybe_unused]] = vector<vector<pair<int,T>>>;
template<typename T> using matrix [[maybe_unused]] = vector<vector<T>>;
const ll INF = 0x3f3f3f3f3f3f3f3f;
/**
* @brief Centroid Tree implementation
*/
class CentroidTree {
int n, root;
int MAX_BIT;
graph<int> g;
vector<int> sz, // size of the subtree
depth, // depth of the node
tree, // centroid tree(tree[i] = parent of i at the centroid tree)
S, T, E; // Euler Tour
vector<ll> dist; // distance from the initial root
int pw = 0;
int dfs(int u, int p) {
sz[u] = 1;
E[++pw] = u;
S[u] = pw;
for (const auto& [v, w]: g[u]) {
if (v == p) continue;
depth[v] = depth[u] + 1;
dist[v] = dist[u] + w;
sz[u] += dfs(v, u);
E[++pw] = u;
}
return sz[u];
}
vector<int> pw2, lg2;
vector<vector<pair<int,int>>> ST;
void lca_prepare(){
pw2.resize(MAX_BIT);
pw2[0] = 1;
for (int i = 1; i < MAX_BIT; i++) pw2[i] = pw2[i - 1] << 1;
lg2.resize(n * 2);
fill(lg2.begin(), lg2.end(), -1);
for (int i = 0; i < MAX_BIT; i++) if (pw2[i] < n * 2) lg2[pw2[i]] = i;
for (int i = 1; i < n * 2; i++) if (lg2[i] == -1) lg2[i] = lg2[i - 1];
ST.resize(MAX_BIT, vector<pair<int,int>>(n * 2));
for (int i = 1; i < n * 2; i++) ST[0][i] = {depth[E[i]], E[i]};
for(int k = 1; k < MAX_BIT; k++) {
for (int i = 1; i < n * 2; i++) {
if (i + pw2[k - 1] > n * 2) continue;
ST[k][i] = min(ST[k - 1][i], ST[k - 1][i + pw2[k - 1]]);
}
}
}
int lca(int u, int v) {
int l = S[u], r = S[v];
if (l > r) swap(l, r);
int k = lg2[r - l + 1];
return min(ST[k][l], ST[k][r - pw2[k] + 1]).second;
}
ll get_dist(int u, int v) {
return dist[u] + dist[v] - 2 * dist[lca(u, v)];
}
int get_centroid(int u) {
for (const auto& [v, w]: g[u]) {
if (sz[u] >> 1 < sz[v] && sz[v] < sz[u]) {
sz[u] -= sz[v];
sz[v] += sz[u];
return get_centroid(v);
}
}
return u;
}
void build_centroid_tree(int u, int p = -1) {
u = get_centroid(u);
if (p == -1) tree[u] = u;
else tree[u] = p;
for (const auto& [v, _]: g[u])
if (sz[v] < sz[u])
build_centroid_tree(v, u);
}
public:
CentroidTree() = default;
explicit
CentroidTree(const graph<int>& g, int root = 1) : g(g), n((int)g.size()), root(root) {
sz.resize(n, 0);
depth.resize(n, 0);
dist.resize(n, 0);
S.resize(n), T.resize(n), E.resize(n * 2);
tree.resize(n);
MAX_BIT = std::ceil(std::log2(n));
dfs(root, -1);
lca_prepare();
build_centroid_tree(root);
x_dist.resize(n, INF);
}
vector<ll> x_dist;
ll ans = INF;
void query(int u, bool x) {
// climb up the centroid tree and update ancestor's set
if (x) {
for (int v = u;; v = tree[v]) {
x_dist[v] = min(x_dist[v], get_dist(u, v));
if (tree[v] == v) break;
}
} else {
for (int v = u;; v = tree[v]) {
ans = min(ans, x_dist[v] + get_dist(u, v));
if (tree[v] == v) break;
}
}
}
void clear(int u) {
for (int v = u;; v = tree[v]) {
x_dist[v] = INF;
if (tree[v] == v) break;
}
ans = INF;
}
}; // class CentroidTree
graph<int> g;
CentroidTree solver;
void Init(int N, int A[], int B[], int D[]) {
g.resize(N + 1);
for (int i = 0; i < N - 1; i++) {
g[A[i]+1].emplace_back(B[i]+1, D[i]);
g[B[i]+1].emplace_back(A[i]+1, D[i]);
}
solver = CentroidTree(g);
}
ll Query(int S, int X[], int T, int Y[]) {
for (int i = 0; i < S; i++) solver.clear(X[i]+1);
for (int i = 0; i < S; i++) solver.query(X[i]+1, true);
for (int i = 0; i < T; i++) solver.query(Y[i]+1, false);
return solver.ans;
}
Compilation message
factories.cpp: In constructor 'CentroidTree::CentroidTree(graph<int>&, int)':
factories.cpp:24:16: warning: 'CentroidTree::g' will be initialized after [-Wreorder]
24 | graph<int> g;
| ^
factories.cpp:22:9: warning: 'int CentroidTree::n' [-Wreorder]
22 | int n, root;
| ^
factories.cpp:106:5: warning: when initialized here [-Wreorder]
106 | CentroidTree(const graph<int>& g, int root = 1) : g(g), n((int)g.size()), root(root) {
| ^~~~~~~~~~~~
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Incorrect |
10 ms |
16988 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Incorrect |
2 ms |
16988 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Incorrect |
10 ms |
16988 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |