/**
____ ____ ____ __________________ ____ ____ ____
||I || ||c || ||e || || || ||M || ||a || ||n ||
||__|| ||__|| ||__|| ||________________|| ||__|| ||__|| ||__||
|/__\| |/__\| |/__\| |/________________\| |/__\| |/__\| |/__\|
*/
#include <iostream>
#include <vector>
#include <map>
#include "race.h"
#define maxn 500005
#define PB push_back
#define X first
#define Y second
typedef long long ll;
typedef std::pair <ll, ll> pll;
const ll INF = 4e18;
std::vector<pll> v[maxn];
int n, k;
ll ans;
// We'll use DFS small-to-large per-node maps (distances measured from the node)
std::map<ll,ll>* mp[maxn];
void dfs_sl(int node, int par)
{
// create map for node: distance 0 -> node-count 1
mp[node] = new std::map<ll,ll>();
mp[node]->emplace(0LL, 1LL);
for (auto &nb : v[node])
{
int to = (int)nb.X;
ll w = nb.Y;
if (to == par) continue;
dfs_sl(to, node);
// ensure mp[node] is the larger map (small-to-large)
if (mp[node]->size() < mp[to]->size())
{
// swap pointers: node's map becomes child's map and vice versa
std::swap(mp[node], mp[to]);
// BUT the map now in mp[node] has distances relative to 'to' not 'node'.
// We must shift all keys in mp[node] by +w and add +1 to values (nodes count),
// because we're turning it into distances from 'node'.
std::map<ll,ll>* shifted = new std::map<ll,ll>();
for (auto &pr : *mp[node])
{
ll newd = pr.first + w;
ll newv = pr.second + 1; // include 'node'
auto it = shifted->find(newd);
if (it == shifted->end()) shifted->emplace(newd, newv);
else it->second = std::min(it->second, newv);
}
// delete old mp[node] contents and replace pointer
delete mp[node];
mp[node] = shifted;
}
else
{
// mp[node] is the larger; we will merge mp[to] into mp[node]
// but need to shift mp[to] keys by +w and values by +1 for queries and insertion
// first check complements (without modifying mp[node])
for (auto &pr : *mp[to])
{
ll dist_from_node = pr.first + w;
ll nodes_cnt = pr.second + 1; // include node
ll need = (ll)k - dist_from_node;
auto it = mp[node]->find(need);
if (it != mp[node]->end())
{
ans = std::min(ans, it->second + nodes_cnt);
}
}
// then insert shifted entries
for (auto &pr : *mp[to])
{
ll newd = pr.first + w;
ll newv = pr.second + 1;
auto it = mp[node]->find(newd);
if (it == mp[node]->end()) mp[node]->emplace(newd, newv);
else it->second = std::min(it->second, newv);
}
// free child's map
delete mp[to];
mp[to] = nullptr;
continue; // next child
}
// If we reached here it means we swapped and mp[node] is the previous mp[to] shifted.
// We still need to merge the remaining (smaller) mp[to] (which now sits in mp[to]) into mp[node].
// Note: after swapping we haven't yet done the query/merge for the (old) smaller child's map:
for (auto &pr : *mp[to])
{
ll dist_from_node = pr.first + w;
ll nodes_cnt = pr.second + 1;
// query existing mp[node] for complement
auto it = mp[node]->find((ll)k - dist_from_node);
if (it != mp[node]->end())
ans = std::min(ans, it->second + nodes_cnt);
}
// now insert them
for (auto &pr : *mp[to])
{
ll newd = pr.first + w;
ll newv = pr.second + 1;
auto it = mp[node]->find(newd);
if (it == mp[node]->end()) mp[node]->emplace(newd, newv);
else it->second = std::min(it->second, newv);
}
// delete child's map
delete mp[to];
mp[to] = nullptr;
}
// also check if there exists an element in mp[node] equal to k (distance from node to some node in its subtree)
auto it = mp[node]->find((ll)k);
if (it != mp[node]->end())
ans = std::min(ans, it->second);
}
int best_path(int N, int K, int h[][2], int L[])
{
n = N;
k = K;
ans = INF;
// clear adjacency
for (int i = 1; i <= n; ++i) v[i].clear();
for (int i = 0; i < n - 1; ++i)
{
// input nodes in h[][] are 0-based; original code used +1 indexing
v[h[i][0] + 1].PB({h[i][1] + 1, L[i]});
v[h[i][1] + 1].PB({h[i][0] + 1, L[i]});
}
for (int i = 0; i <= n; ++i) mp[i] = nullptr;
dfs_sl(1, 0);
// free top-level map
if (mp[1]) { delete mp[1]; mp[1] = nullptr; }
if (ans == INF) return -1;
else return (int)(ans - 1); // stored values are node counts (edges+1)
}
| # | 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... |