#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>
#include <map>
#include <vector> // Add missing include
// --- Disjoint Set Union (DSU) ---
struct DSU {
std::vector<int> parent;
DSU(int n) {
parent.resize(n + 1);
std::iota(parent.begin(), parent.end(), 0);
}
int find(int x) {
if (parent[x] == x) return x;
return parent[x] = find(parent[x]);
}
bool unite(int x, int y) {
int rootX = find(x);
int rootY = find(y);
if (rootX != rootY) {
parent[rootX] = rootY;
return true;
}
return false;
}
};
// --- Edge Structures ---
struct Edge {
int u, v, cost, original_id; // original_id is 0 to M-1
};
struct NewEdgeInput {
int id; // 0 to K-1
int u, v;
};
// --- MST T0 & LCA / Path Max ---
const int MAXN = 100005;
const int LOGN = 18; // Max ~log2(MAXN)
// Structure to store MST edges with neighbor, cost, and original_id
struct MSTAdjInfo {
int neighbor;
int cost;
int original_id;
};
std::vector<MSTAdjInfo> adj_mst[MAXN]; // Adjacency list for MST T0
// For Binary Lifting LCA/Path Max
int up[MAXN][LOGN];
int max_edge_cost[MAXN][LOGN]; // Stores max cost on path
int max_edge_id[MAXN][LOGN]; // Stores original_id of the max cost edge
int depth[MAXN];
std::map<int, long long> edge_flow; // Map original_id to its flow in T0
std::vector<long long> p_counts; // Participant counts (0-indexed)
// DFS for LCA, Path Max precomputation, and storing parent edge info
void dfs_lca_build(int u, int p_node, int p_cost, int p_edge_id, int d) {
depth[u] = d;
up[u][0] = p_node;
max_edge_cost[u][0] = p_cost;
max_edge_id[u][0] = p_edge_id;
for (int i = 1; i < LOGN; ++i) {
int ancestor = up[u][i - 1];
if (ancestor != 0) { // Check if ancestor exists
up[u][i] = up[ancestor][i - 1];
if (max_edge_cost[u][i - 1] >= max_edge_cost[ancestor][i - 1]) {
max_edge_cost[u][i] = max_edge_cost[u][i - 1];
max_edge_id[u][i] = max_edge_id[u][i - 1];
} else {
max_edge_cost[u][i] = max_edge_cost[ancestor][i - 1];
max_edge_id[u][i] = max_edge_id[ancestor][i - 1];
}
} else { // Handle case where 2^i ancestor doesn't exist
up[u][i] = 0;
max_edge_cost[u][i] = max_edge_cost[u][i-1]; // Inherit from previous level
max_edge_id[u][i] = max_edge_id[u][i-1];
}
}
for (const auto& edge_info : adj_mst[u]) {
int v = edge_info.neighbor;
if (v != p_node) {
dfs_lca_build(v, u, edge_info.cost, edge_info.original_id, d + 1);
}
}
}
// DFS to calculate subtree participant sums and edge flows
long long dfs_flow_calc(int u, int p_node) {
long long current_subtree_p = p_counts[u - 1]; // p_counts is 0-indexed
for (const auto& edge_info : adj_mst[u]) {
int v = edge_info.neighbor;
if (v != p_node) {
long long child_flow = dfs_flow_calc(v, u);
edge_flow[edge_info.original_id] = child_flow; // Store flow by original_id
current_subtree_p += child_flow;
}
}
return current_subtree_p;
}
// Query Max Edge Info on Path u-v
std::pair<int, int> query_path_max(int u, int v) {
int max_c = 0;
int max_id = -1;
if (depth[u] < depth[v]) std::swap(u, v);
for (int i = LOGN - 1; i >= 0; --i) {
if (up[u][i] != 0 && depth[u] - (1 << i) >= depth[v]) {
if (max_edge_cost[u][i] > max_c) {
max_c = max_edge_cost[u][i];
max_id = max_edge_id[u][i];
} else if (max_edge_cost[u][i] == max_c) {
// Tie-breaking? Assume distinct costs for now, or handle if needed.
}
u = up[u][i];
}
}
if (u == v) return {max_c, max_id};
for (int i = LOGN - 1; i >= 0; --i) {
if (up[u][i] != up[v][i]) {
// Compare and update max_c, max_id for both u and v paths
if (max_edge_cost[u][i] > max_c) { max_c = max_edge_cost[u][i]; max_id = max_edge_id[u][i]; }
if (max_edge_cost[v][i] > max_c) { max_c = max_edge_cost[v][i]; max_id = max_edge_id[v][i]; }
u = up[u][i];
v = up[v][i];
}
}
// Compare last edges to LCA
if (max_edge_cost[u][0] > max_c) { max_c = max_edge_cost[u][0]; max_id = max_edge_id[u][0]; }
if (max_edge_cost[v][0] > max_c) { max_c = max_edge_cost[v][0]; max_id = max_edge_id[v][0]; }
return {max_c, max_id}; // {max_cost, original_id_of_max_cost_edge}
}
// --- MWIS Solver (Backtracking) ---
std::vector<std::vector<int>> conflict_adj; // Conflict graph adjacency list
std::vector<long long> node_weights; // Revenue R_i for node i (0 to K-1)
long long max_revenue_found = 0;
std::vector<bool> is_selected; // Current selection in backtracking
void solve_mwis(int k_idx, long long current_revenue) {
if (k_idx == node_weights.size()) {
max_revenue_found = std::max(max_revenue_found, current_revenue);
return;
}
// Option 1: Don't include node k_idx
solve_mwis(k_idx + 1, current_revenue);
// Option 2: Try to include node k_idx
bool conflict_exists = false;
for (int neighbor : conflict_adj[k_idx]) {
if (neighbor < k_idx && is_selected[neighbor]) {
conflict_exists = true;
break;
}
}
if (!conflict_exists) {
is_selected[k_idx] = true;
solve_mwis(k_idx + 1, current_revenue + node_weights[k_idx]);
is_selected[k_idx] = false; // Backtrack
}
}
// --- Main ---
int main() {
std::ios_base::sync_with_stdio(false);
std::cin.tie(NULL);
int n, m, k_val;
std::cin >> n >> m >> k_val;
std::vector<Edge> original_edges_input(m);
for (int i = 0; i < m; ++i) {
std::cin >> original_edges_input[i].u >> original_edges_input[i].v >> original_edges_input[i].cost;
original_edges_input[i].original_id = i;
}
std::vector<NewEdgeInput> new_edges_input(k_val);
for (int i = 0; i < k_val; ++i) {
new_edges_input[i].id = i;
std::cin >> new_edges_input[i].u >> new_edges_input[i].v;
}
p_counts.resize(n);
for (int i = 0; i < n; ++i) {
std::cin >> p_counts[i];
}
// 1. Compute MST T0
std::sort(original_edges_input.begin(), original_edges_input.end(), [](const Edge& a, const Edge& b) {
return a.cost < b.cost;
});
DSU dsu_mst(n);
int mst_edge_count = 0;
for (const auto& edge : original_edges_input) {
if (dsu_mst.unite(edge.u, edge.v)) {
adj_mst[edge.u].push_back({edge.v, edge.cost, edge.original_id});
adj_mst[edge.v].push_back({edge.u, edge.cost, edge.original_id});
mst_edge_count++;
if (mst_edge_count == n - 1) break;
}
}
// 2. Precompute LCA/Path Max and Flows
dfs_lca_build(1, 0, 0, -1, 0); // Build LCA/Path Max from root 1
dfs_flow_calc(1, 0); // Calculate flows from root 1
// 3. Calculate potential revenue and conflicts
node_weights.resize(k_val, 0);
std::map<int, std::vector<int>> replaced_to_replacers; // original_id -> list of new_edge ids (0..k-1)
for (int i = 0; i < k_val; ++i) {
int u = new_edges_input[i].u;
int v = new_edges_input[i].v;
std::pair<int, int> path_max_info = query_path_max(u, v);
int max_c = path_max_info.first;
int replaced_original_id = path_max_info.second;
if (replaced_original_id != -1 && edge_flow.count(replaced_original_id)) {
long long revenue = (long long)max_c * edge_flow[replaced_original_id];
node_weights[i] = revenue;
replaced_to_replacers[replaced_original_id].push_back(i);
}
// Else, revenue remains 0
}
// 4. Build conflict graph
conflict_adj.resize(k_val);
for (auto const& [replaced_id, replacer_indices] : replaced_to_replacers) {
if (replacer_indices.size() > 1) {
for (size_t i = 0; i < replacer_indices.size(); ++i) {
for (size_t j = i + 1; j < replacer_indices.size(); ++j) {
int u_idx = replacer_indices[i];
int v_idx = replacer_indices[j];
conflict_adj[u_idx].push_back(v_idx);
conflict_adj[v_idx].push_back(u_idx);
}
}
}
}
// 5. Solve MWIS
is_selected.assign(k_val, false);
max_revenue_found = 0;
solve_mwis(0, 0);
// 6. Output result
std::cout << max_revenue_found << std::endl;
return 0;
}
# | 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... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |