#include "closing.h"
#include <bits/stdc++.h>
typedef std::int64_t i64;
typedef std::size_t usize;
typedef std::ptrdiff_t isize;
struct edge
{
usize dest;
i64 length;
};
struct node
{
std::vector<edge> edges;
i64 distance_x = std::numeric_limits<i64>::max();
i64 distance_y = std::numeric_limits<i64>::max();
i64 &distance(bool id)
{
if (id)
return distance_y;
else
return distance_x;
}
};
void mark_distances(std::vector<node> &nodes, usize start, bool id)
{
std::queue<std::pair<i64, usize>> queue;
queue.push({0, start});
while (!queue.empty())
{
i64 dist = queue.front().first;
usize n = queue.front().second;
queue.pop();
nodes[n].distance(id) = dist;
for (edge e : nodes[n].edges)
{
i64 new_dist = dist + e.length;
if (nodes[e.dest].distance(id) > dist)
queue.push({new_dist, e.dest});
}
}
}
inline i64 sat_add(i64 a, i64 b)
{
if (std::numeric_limits<i64>::max() - b < a)
return std::numeric_limits<i64>::max();
return a + b;
}
struct poss_conn
{
i64 cost;
usize city;
bool channel;
};
bool operator<(poss_conn const &a, poss_conn const &b)
{
return a.cost > b.cost;
}
struct graph
{
std::vector<node> nodes;
std::vector<bool> conn_x;
std::vector<bool> conn_y;
usize count;
i64 cost;
i64 max_cost;
std::priority_queue<poss_conn> queue;
graph(std::vector<node> nodes, i64 max_cost) : nodes(nodes), conn_x(nodes.size(), false), conn_y(nodes.size(), false), count(0), cost(0), max_cost(max_cost) {}
std::vector<bool> &conns(bool channel)
{
return channel ? conn_y : conn_x;
}
i64 cost_of(usize id, bool channel)
{
i64 cost = nodes[id].distance(channel);
if (conns(!channel)[id])
cost -= nodes[id].distance(!channel);
return std::max(cost, (i64)0);
}
void mark_for_sale(usize id, bool channel)
{
queue.push({cost_of(id, channel), id, channel});
}
bool cost_ok()
{
return cost <= max_cost;
}
bool buy(usize id, bool channel)
{
if (conns(channel)[id])
return true;
cost += cost_of(id, channel);
if (!cost_ok())
return false;
conns(channel)[id] = true;
count++;
if (!conns(!channel)[id])
{
bool reachable = false;
for (edge &e : nodes[id].edges)
if (conns(!channel)[e.dest])
reachable = true;
if (reachable)
mark_for_sale(id, !channel);
}
for (edge &e : nodes[id].edges)
if (!conns(channel)[e.dest])
mark_for_sale(e.dest, channel);
return true;
}
bool run_queue()
{
if (queue.empty())
return false;
poss_conn conn = queue.top();
queue.pop();
return buy(conn.city, conn.channel);
}
};
i64 get_best_score(std::vector<node> &nodes, i64 max_cost, usize city_x, usize city_y)
{
usize best_count = 0;
for (usize prebuy = city_x; prebuy < nodes.size(); prebuy++)
{
graph g(nodes, max_cost);
for (usize n = city_x; n <= prebuy; n++)
{
g.buy(n, false);
}
g.buy(city_y, true);
if (!g.cost_ok())
continue;
while (g.run_queue())
;
if (g.count > best_count)
best_count = g.count;
}
return best_count;
}
int max_score(int city_count, int city_x, int city_y, long long max_closing_sum, std::vector<int> starts, std::vector<int> ends, std::vector<int> lengths)
{
if (city_y < city_x)
std::swap(city_x, city_y);
std::vector<node> nodes(city_count);
for (usize i = 0; i < city_count - 1; i++)
{
nodes[starts[i]].edges.push_back(edge{(usize)ends[i], lengths[i]});
nodes[ends[i]].edges.push_back(edge{(usize)starts[i], lengths[i]});
}
mark_distances(nodes, city_x, false);
mark_distances(nodes, city_y, true);
return get_best_score(nodes, max_closing_sum, city_x, city_y);
}
