P2212 [USACO14MAR] Watering the Fields S 最小生成树prim 堆优化

解题思路与代码注释

题目理解

这道题目要求我们在给定的N个点之间建立灌溉系统，使得所有点都能连通。建造成本为两点间欧几里得距离的平方，且只有当成本≥C时才能建造管道。我们需要找到使所有点连通的最小总成本，如果无法满足条件则输出-1。

解题方法

使用Prim算法来求解最小生成树，并在过程中：

1. 仅考虑距离平方≥C的边

2. 计算满足条件的最小生成树总成本

3. 检查是否所有点都被连通

温馨提示：本题无需开根号，但笔者习惯写了double类型，理论上int类型也可以通过

代码详细注释

#include<bits/stdc++.h>
#define pii pair<int,int>  // 定义pair类型，用于优先队列存储(距离,点)
using namespace std;

const int N = 2e5 + 10, inf = 0x3f3f3f3f;

struct node {
    int to;        // 边的终点
    double val;    // 边的权值(距离平方)
};

int n, c, x[N], y[N];  // n-点数，c-最小成本限制，x/y-坐标数组
vector<node> g[N];     // 邻接表存储图
int dis[N], vis[N];    // dis-各点到生成树的最小距离，vis-标记是否已加入生成树

// 计算两点间距离平方
double dist(int i, int j) {
    double a = x[i] - x[j];
    double b = y[i] - y[j];
    return a * a + b * b;
}

void prim() {
    // 小根堆优化Prim算法
    priority_queue<pii, vector<pii>, greater<pii>> q;
    memset(dis, inf, sizeof(dis));  // 初始化距离为无穷大
    dis[1] = 0;                     // 从1号点开始
    q.push({0, 1});                 // 将起点加入优先队列
    
    double sum = 0, ans = 0;        // sum-已连接点数，ans-总成本
    
    while(q.size()) {
        int x = q.top().second;     // 当前距离生成树最近的点
        q.pop();
        
        if(vis[x]) continue;        // 已加入则跳过
        vis[x] = 1;                 // 标记为已加入
        ans += dis[x];              // 累加边权
        sum++;                      // 计数
        
        // 遍历所有邻接边
        for(int i = 0; i < g[x].size(); i++) {
            int y = g[x][i].to;     // 邻接点
            int z = g[x][i].val;    // 边权
            if(!vis[y] && dis[y] > z) {  // 如果y未加入且找到更小距离
                dis[y] = z;         // 更新距离
                q.push({dis[y], y});// 加入优先队列
            }
        }
    }
    
    // 检查是否所有点都连通且总成本≥c
    if(sum != n || ans < c) cout << -1;
    else cout << (int)ans;          // 输出整数结果
}

int main() {
    cin >> n >> c;
    // 读入所有点坐标
    for(int i = 1; i <= n; i++) 
        cin >> x[i] >> y[i];
    
    // 构建邻接表，只保留距离平方≥c的边
    for(int i = 1; i <= n; i++) {
        for(int j = i + 1; j <= n; j++) {
            double val = dist(i, j);
            if(val < c) continue;   // 跳过不满足条件的边
            g[i].push_back({j, val});
            g[j].push_back({i, val});
        }
    }
    
    prim();  // 执行Prim算法
    return 0;
}