#include<bits/stdc++.h>
#include "factories.h"
using namespace std;
#define int long long
int n, a[500005], b[500005], d[500005], tin[500005], tout[500005], timer = -1, s, t, x[500005], y[500005], q, w[500005], dp[2][500005], col[500005];
vector<pair<int,int>> adj[500005], adj2[500005];
pair<int, int> up[20][500005];
vector<int> euler;
int res = 1e18;
bool cmp(int a, int b) {
return tin[a] < tin[b];
}
void dfs(int u, int par) {
tin[u] = ++timer;
euler.push_back(u);
for (pair<int, int> node: adj[u]) {
if(node.first == par) continue;
w[node.first] = w[u] + node.second;
dfs(node.first, u);
timer++;
euler.push_back(u);
}
tout[u] = timer;
}
void build() {
for (int i = 0; i < euler.size(); i++) up[0][i] = {tin[euler[i]], euler[i]};
for (int i = 1; i < 20; i++) {
for (int j = 0; j + (1<<i) - 1 < euler.size(); j++) {
up[i][j] = min(up[i-1][j], up[i-1][j+(1<<(i-1))]);
}
}
}
pair<int, int> get(int u, int v) {
int lg = __lg(v - u + 1);
return min(up[lg][u], up[lg][v-(1<<lg)+1]);
}
void dfs2(int u, int par) {
dp[0][u] = dp[1][u] = 1e18;
for (pair<int, int> node: adj2[u]) {
if(node.first == par) continue;
dfs2(node.first, u);
dp[0][u] = min(dp[0][u], dp[0][node.first] + node.second);
dp[1][u] = min(dp[1][u], dp[1][node.first] + node.second);
if(col[node.first] == 1) dp[0][u] = min(dp[0][u], w[node.first] - w[u]);
if(col[node.first] == 2) dp[1][u] = min(dp[1][u], w[node.first] - w[u]);
}
if(col[u] == 1) res = min(res, dp[1][u]);
if(col[u] == 2) res = min(res, dp[0][u]);
res = min(res, dp[1][u] + dp[0][u]);
}
long long query(int s, int x[], int t, int y[]) {
vector<int> v, v2;
int root, mx = 1e18;
for (int i = 1; i <= s; i++) {x[i]++, col[x[i]] = 1, v.push_back(x[i]); if(mx > tin[x[i]]) mx = tin[x[i]], root = x[i];}
for (int i = 1; i <= t; i++) {y[i]++, col[y[i]] = 2, v.push_back(y[i]); if(mx > tin[y[i]]) mx = tin[y[i]], root = y[i];}
sort(v.begin(), v.end(), cmp);
v2 = v;
for (int i = 0; i < v.size() - 1; i++) {
pair<int, int> lca = get(tin[v[i]], tin[v[i+1]]);
v2.push_back(lca.second);
if(mx > lca.first) mx = lca.first, root = lca.second;
}
sort(v2.begin(), v2.end(), cmp);
stack<int> S;
S.push(v2[0]);
for (int i = 1; i < v2.size(); i++) {
int cur = S.top(), p = v2[i];
if(tin[p] == tin[cur]) continue;
while(tin[cur] > tin[p] || tout[cur] < tin[p]) {
S.pop();
cur = S.top();
}
adj2[cur].push_back({p, w[p] - w[cur]});
S.push(p);
}
res = 1e18;
dfs2(root, -1);
for (int i = 1; i <= s; i++) col[x[i]] = 0;
for (int i = 1; i <= t; i++) col[y[i]] = 0;
return res;
}
void init(int n, int a[], int b[], int d[]) {
cin >> q;
for (int i = 1; i < n; i++) {
a[i]++, b[i]++;
adj[a[i]].push_back({b[i], d[i]});
adj[b[i]].push_back({a[i], d[i]});
}
dfs(1, -1);
build();
while(q--) {
cout << Query() << '\n';
}
}
Compilation message
factories.cpp: In function 'void build()':
factories.cpp:27:20: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
27 | for (int i = 0; i < euler.size(); i++) up[0][i] = {tin[euler[i]], euler[i]};
| ~~^~~~~~~~~~~~~~
factories.cpp:29:34: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
29 | for (int j = 0; j + (1<<i) - 1 < euler.size(); j++) {
| ~~~~~~~~~~~~~~~^~~~~~~~~~~~~~
factories.cpp: In function 'long long int query(long long int, long long int*, long long int, long long int*)':
factories.cpp:59:20: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
59 | for (int i = 0; i < v.size() - 1; i++) {
| ~~^~~~~~~~~~~~~~
factories.cpp:67:20: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
67 | for (int i = 1; i < v2.size(); i++) {
| ~~^~~~~~~~~~~
factories.cpp: In function 'void init(long long int, long long int*, long long int*, long long int*)':
factories.cpp:93:17: error: too few arguments to function 'long long int Query(int, int*, int, int*)'
93 | cout << Query() << '\n';
| ^
In file included from factories.cpp:2:
factories.h:5:11: note: declared here
5 | long long Query(int S, int X[], int T, int Y[]);
| ^~~~~
