#include "race.h"
#include <iostream>
#include <vector>
#include <unordered_set>
int INF = 1000000000;
int k;
std::vector<int> subtreesize;
std::vector<std::pair<int, int>> possible;
void dfs(std::vector<std::vector<std::pair<int, int>>>& adj, int node, int p, std::unordered_set<int>& used) {
subtreesize[node] = 0;
if (used.count(node)) return;
subtreesize[node] = 1;
for (auto [to, _d] : adj[node]) {
if (to == p) continue;
dfs(adj, to, node, used);
subtreesize[node] += subtreesize[to];
}
}
// subsolve is called once for each node for each layer
void subsolve(int& ans, std::vector<std::vector<std::pair<int, int>>>& adj, int node, int respectivecentroid, int d, int steps, int p, std::vector<std::pair<int, int>>& subpossible, std::unordered_set<int>& used) {
if (used.count(node)) return;
if (d > k) return;
subpossible.emplace_back(d, steps);
if (possible[k-d].second == respectivecentroid) {
ans = std::min(ans, steps + possible[k-d].first);
}
for (auto [to, dd] : adj[node]) {
if (to == p) continue;
subsolve(ans, adj, to, respectivecentroid, d+dd, steps+1, node, subpossible, used);
}
}
// solve is called once for each centroid
int solve(std::vector<std::vector<std::pair<int, int>>>& adj, int node, std::unordered_set<int>& used) {
// std::cout << node << std::endl;
int ans = INF;
possible[0] = {0, node};
for (auto [to, d] : adj[node]) {
if (used.count(to)) continue;
std::vector<std::pair<int, int>> subpossible;
subsolve(ans, adj, to, node, d, 1, node, subpossible, used);
for (auto [w, u] : subpossible) {
if (possible[w].second == node) {
possible[w].first = std::min(possible[w].first, u);
} else {
possible[w] = std::make_pair(u, node);
}
}
}
return ans;
}
int getcentroid(std::vector<std::vector<std::pair<int, int>>>& adj, int node, std::unordered_set<int>& used) {
dfs(adj, node, -1, used);
int n = subtreesize[node];
int cur = node;
while (true) {
bool found = false;
for (auto [to, d] : adj[cur]) {
if (used.count(to)) continue;
if (subtreesize[to] > subtreesize[cur]) continue; // parent
if (subtreesize[to]*2 >= n) {
cur = to;
found = true;
break;
}
}
if (!found) break;
}
return cur;
}
// bigsolve is called once for each component
int bigsolve(std::vector<std::vector<std::pair<int, int>>>& adj, int node, std::unordered_set<int>& used) {
int ans = INF;
node = getcentroid(adj, node, used);
used.insert(node);
ans = std::min(ans, solve(adj, node, used));
for (auto [to, d] : adj[node]) {
if (used.count(to)) continue;
ans = std::min(ans, bigsolve(adj, to, used));
}
return ans;
}
int best_path(int N, int K, int H[][2], int L[])
{
k = K;
subtreesize.resize(N);
possible.assign(K+1, std::make_pair(-1, -1));
std::vector<std::vector<std::pair<int, int>>> adj(N);
for (int i = 0; i < N-1; ++i) {
int u, v, l;
u = H[i][0], v = H[i][1], l = L[i];
adj[u].emplace_back(v, l);
adj[v].emplace_back(u, l);
}
std::unordered_set<int> used;
int min = bigsolve(adj, 0, used);
if (min == INF) min = -1;
return min;
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |