This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#pragma GCC optimize("O3")
#include <factories.h>
#include <iostream>
#include <algorithm>
#include <vector>
#include <cassert>
#include <cstring>
#include <numeric>
#include <set>
#include <queue>
#define sz(a) (int)((a).size())
#define all(a) (a).begin(), (a).end()
#define pb push_back
using namespace std;
using ll = long long;
using vi = vector<int>;
using pii = pair<int, int>;
const int LG = 20;
const int N = 1000228;
const ll inf = 1e18;
int n;
vector<vector<pii>> g;
vi in, out;
int timer;
vector<ll> dep;
vector<bool> in_A1, in_A2;
int lg[N];
pair<ll, int> sparse[LG][N];
vi A;
vector<vector<pair<int, ll>>> ga;
vector<ll> dp;
ll ans;
void dfs(int v, int p) {
sparse[0][timer] = { dep[v], v };
in[v] = timer++;
out[v] = in[v];
for (auto e : g[v]) {
if (e.first != p) {
dep[e.first] = dep[v] + e.second;
dfs(e.first, v);
sparse[0][timer] = { dep[v], v };
out[v] = timer++;
}
}
}
void Init(int _n, int _u[], int _v[], int _w[]) {
n = _n;
g.resize(n);
for (int i = 0; i < n - 1; i++) {
g[_u[i]].pb({ _v[i], _w[i] });
g[_v[i]].pb({ _u[i], _w[i] });
}
dep.resize(n);
in.resize(n);
out.resize(n);
dfs(0, -1);
for (int i = 2; i < N; i++) {
lg[i] = 1 + lg[i >> 1];
}
for (int j = 1; j < LG; j++) {
for (int i = 0; i + (1 << j) - 1 < timer; i++) {
sparse[j][i] = min(sparse[j - 1][i], sparse[j - 1][i + (1 << (j - 1))]);
}
}
in_A1.resize(n, false);
in_A2.resize(n, false);
}
bool upper(int u, int v) {
return (in[u] <= in[v] && out[v] <= out[u]);
}
int LCA(int u, int v) {
u = out[u];
v = out[v];
if (u > v) {
swap(u, v);
}
int l = lg[v - u + 1];
return min(sparse[l][u], sparse[l][v - (1 << l) + 1]).second;
}
void dfs_dp(int v) {
if (in_A2[A[v]]) {
dp[v] = 0;
}
for (auto e : ga[v]) {
dfs_dp(e.first);
dp[v] = min(dp[v], dp[e.first] + e.second);
}
}
void dfs_upd(int v, ll dp_up) {
if (in_A1[A[v]]) {
ans = min(ans, min(dp[v], dp_up));
}
vector<ll> suf_min(sz(ga[v]) + 1);
suf_min.back() = inf;
for (int i = sz(ga[v]) - 1; i >= 0; i--) {
suf_min[i] = min(suf_min[i + 1], dp[ga[v][i].first] + ga[v][i].second);
}
ll pref_min = inf;
for (int i = 0; i < sz(ga[v]); i++) {
dfs_upd(ga[v][i].first, min({ pref_min, suf_min[i + 1], dp_up, in_A2[A[v]] ? 0 : inf }) + ga[v][i].second);
pref_min = min(pref_min, dp[ga[v][i].first] + ga[v][i].second);
}
}
ll Query(int n1, int _A1[], int n2, int _A2[]) {
A.clear();
A.resize(n1 + n2);
for (int i = 0; i < n1; i++) {
A[i] = _A1[i];
in_A1[_A1[i]] = true;
}
for (int i = 0; i < n2; i++) {
A[i + n1] = _A2[i];
in_A2[_A2[i]] = true;
}
sort(all(A), [](int u, int v) { return in[u] < in[v]; });
for (int i = sz(A) - 2; i >= 0; i--) {
A.pb(LCA(A[i], A[i + 1]));
}
sort(all(A), [](int u, int v) { return in[u] < in[v]; });
A.resize(unique(all(A)) - A.begin());
ga.clear();
ga.resize(sz(A));
vi st;
int edge_check = 0;
for (int i = 0; i < sz(A); i++) {
while (!st.empty() && !upper(A[st.back()], A[i])) {
st.pop_back();
}
if (!st.empty()) {
ga[st.back()].pb({ i, dep[A[i]] - dep[A[st.back()]] });
edge_check++;
}
st.pb(i);
}
assert(edge_check == sz(A) - 1);
dp.assign(sz(A), inf);
dfs_dp(0);
ans = inf;
dfs_upd(0, inf);
for (int i = 0; i < n1; i++) {
in_A1[_A1[i]] = false;
}
for (int i = 0; i < n2; i++) {
in_A2[_A2[i]] = false;
}
return ans;
}
/*
int u[100], v[100], w[100];
int A1[100], A2[100];
int32_t main() {
ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
int n, q;
cin >> n >> q;
for (int i = 0; i < n - 1; i++) {
cin >> u[i] >> v[i] >> w[i];
}
Init(n, u, v, w);
while (q--) {
int n1, n2;
cin >> n1 >> n2;
for (int i = 0; i < n1; i++) {
cin >> A1[i];
}
for (int i = 0; i < n2; i++) {
cin >> A2[i];
}
cout << Query(n1, A1, n2, A2) << endl;
}
}
*/
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |