#include <iostream>
#include <vector>
#include <algorithm>
#include <numeric>
using namespace std;
struct Edge {
int u, v;
long long cost;
bool is_new; // To distinguish new road from existing ones
};
// Utility functions for the Union-Find data structure
int find(vector<int>& parent, int i) {
return parent[i] == i ? i : parent[i] = find(parent, parent[i]);
}
bool union_set(vector<int>& parent, vector<int>& rank, int a, int b) {
a = find(parent, a);
b = find(parent, b);
if (a != b) {
if (rank[a] < rank[b]) swap(a, b);
parent[b] = a;
if (rank[a] == rank[b]) rank[a]++;
return true;
}
return false;
}
// Function to compute MST using Kruskal's algorithm and return the maximum revenue
long long compute_mst_and_max_revenue(int N, vector<Edge>& edges, int x, int y, const vector<int>& p) {
sort(edges.begin(), edges.end(), [](const Edge& a, const Edge& b) {
return a.cost < b.cost;
});
vector<int> parent(N + 1), rank(N + 1, 0);
iota(parent.begin(), parent.end(), 0);
long long max_revenue = 0;
vector<int> subtree_size(N + 1, 0);
vector<bool> in_mst(N - 1, false);
// Compute MST and participant traffic without considering new road
int edges_used = 0;
for (size_t i = 0; i < edges.size() && edges_used < N - 1; ++i) {
if (union_set(parent, rank, edges[i].u, edges[i].v)) {
in_mst[i] = true;
edges_used++;
}
}
// Calculate participants passing through each edge in MST
function<void(int, int)> dfs = [&](int node, int parent) {
subtree_size[node] = p[node];
for (const auto& e : edges) {
if (in_mst[&e - &edges[0]] && ((e.u == node && e.v != parent) || (e.v == node && e.u != parent))) {
dfs(e.u == node ? e.v : e.u, node);
subtree_size[node] += subtree_size[e.u == node ? e.v : e.u];
}
}
};
dfs(1, -1);
// Try different toll settings for the new road
for (long long toll = 1; toll <= 1000000; toll++) {
Edge new_road = {x, y, toll, true};
edges.push_back(new_road);
iota(parent.begin(), parent.end(), 0);
fill(rank.begin(), rank.end(), 0);
long long revenue = 0;
edges_used = 0;
// Recompute MST including new road
for (size_t i = 0; i < edges.size() && edges_used < N; ++i) {
if (union_set(parent, rank, edges[i].u, edges[i].v)) {
if (edges[i].is_new) {
revenue = toll * (subtree_size[x] + subtree_size[y]);
}
edges_used++;
}
}
max_revenue = max(max_revenue, revenue);
edges.pop_back(); // Remove new road for next iteration
}
return max_revenue;
}
int main() {
int N, M, K;
cin >> N >> M >> K;
vector<Edge> edges(M);
for (int i = 0; i < M; i++) {
cin >> edges[i].u >> edges[i].v >> edges[i].cost;
edges[i].is_new = false;
}
int x, y;
cin >> x >> y; // Read the new road endpoints
vector<int> p(N + 1);
for (int i = 1; i <= N; i++) {
cin >> p[i];
}
// Calculate the maximum revenue
long long max_revenue = compute_mst_and_max_revenue(N, edges, x, y, p);
cout << max_revenue << endl;
return 0;
}
Compilation message
toll.cpp: In function 'long long int compute_mst_and_max_revenue(int, std::vector<Edge>&, int, int, const std::vector<int>&)':
toll.cpp:54:5: error: 'function' was not declared in this scope
54 | function<void(int, int)> dfs = [&](int node, int parent) {
| ^~~~~~~~
toll.cpp:5:1: note: 'std::function' is defined in header '<functional>'; did you forget to '#include <functional>'?
4 | #include <numeric>
+++ |+#include <functional>
5 |
toll.cpp:54:27: error: expression list treated as compound expression in functional cast [-fpermissive]
54 | function<void(int, int)> dfs = [&](int node, int parent) {
| ^
toll.cpp:54:14: error: expected primary-expression before 'void'
54 | function<void(int, int)> dfs = [&](int node, int parent) {
| ^~~~
toll.cpp:63:5: error: 'dfs' was not declared in this scope
63 | dfs(1, -1);
| ^~~
