自定义小根堆

在C++中,priority_queue 默认是一个 大根堆(即最大元素在堆顶)。如果想将其改为 小根堆(即最小元素在堆顶),可以通过以下两种方式实现:

方法 1:使用 greater 比较函数

priority_queue 的第三个模板参数可以指定比较函数。使用 greater 可以让 priority_queue 成为小根堆。

修改后的代码:

#include <iostream>
#include <vector>
#include <queue>
#include <cmath>
using namespace std;

struct node1 {
    double x;
    double y;
    int num;
};

struct node {
    node1 x;
    node1 y;
    double dis;
    node(node1 a, node1 b) {
        this->x = a;
        this->y = b;
        dis = sqrt((a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y));
    }
    // 重载 < 运算符,用于比较
    bool operator<(const node& other) const {
        return dis > other.dis;  // 小根堆
    }
};

int parent[101];

void init() {
    for (int i = 0; i < 101; i++) {
        parent[i] = i;
    }
}

int find(int x) {
    if (parent[x] != x) {
        parent[x] = find(parent[x]);
    }
    return parent[x];
}

void Union(int x, int y) {
    int roota = find(x);
    int rootb = find(y);
    if (roota != rootb) {
        parent[roota] = rootb;
    }
}

int main() {
    int n;
    cin >> n;
    vector<node1> res;
    init();
    for (int i = 0; i < n; i++) {
        double x, y;
        cin >> x >> y;
        node1 d;
        d.x = x; d.y = y; d.num = i;
        res.push_back(d);
    }

    // 使用 greater 比较函数,使 priority_queue 成为小根堆
    priority_queue<node, vector<node>, greater<node>> sss;
    for (int i = 0; i < res.size(); i++) {
        for (int j = i + 1; j < res.size(); j++) {
            sss.push(node(res[i], res[j]));
        }
    }

    double sum = 0;
    while (!sss.empty()) {
        node a = sss.top(); sss.pop();
        if (find(a.x.num) != find(a.y.num)) {
            sum += a.dis;
            Union(a.x.num, a.y.num);
        }
    }

    printf("%.2f\n", sum);  // 保留两位小数
}

方法 2:自定义比较函数

如果你不想重载 < 运算符,可以通过自定义比较函数来实现小根堆。

修改后的代码:

#include <iostream>
#include <vector>
#include <queue>
#include <cmath>
using namespace std;

struct node1 {
    double x;
    double y;
    int num;
};

struct node {
    node1 x;
    node1 y;
    double dis;
    node(node1 a, node1 b) {
        this->x = a;
        this->y = b;
        dis = sqrt((a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y));
    }
};

// 自定义比较函数
struct Compare {
    bool operator()(const node& a, const node& b) {
        return a.dis > b.dis;  // 小根堆
    }
};

int parent[101];

void init() {
    for (int i = 0; i < 101; i++) {
        parent[i] = i;
    }
}

int find(int x) {
    if (parent[x] != x) {
        parent[x] = find(parent[x]);
    }
    return parent[x];
}

void Union(int x, int y) {
    int roota = find(x);
    int rootb = find(y);
    if (roota != rootb) {
        parent[roota] = rootb;
    }
}

int main() {
    int n;
    cin >> n;
    vector<node1> res;
    init();
    for (int i = 0; i < n; i++) {
        double x, y;
        cin >> x >> y;
        node1 d;
        d.x = x; d.y = y; d.num = i;
        res.push_back(d);
    }

    // 使用自定义比较函数,使 priority_queue 成为小根堆
    priority_queue<node, vector<node>, Compare> sss;
    for (int i = 0; i < res.size(); i++) {
        for (int j = i + 1; j < res.size(); j++) {
            sss.push(node(res[i], res[j]));
        }
    }

    double sum = 0;
    while (!sss.empty()) {
        node a = sss.top(); sss.pop();
        if (find(a.x.num) != find(a.y.num)) {
            sum += a.dis;
            Union(a.x.num, a.y.num);
        }
    }

    printf("%.2f\n", sum);  // 保留两位小数
}

关键点总结

  1. priority_queue 默认是大根堆,可以通过 greater 或自定义比较函数改为小根堆。
  2. 方法 1:使用 greater<node>,但需要重载 < 运算符。
  3. 方法 2:自定义比较函数,不需要重载 < 运算符。
  4. 两种方法都可以实现小根堆,选择适合你的方式即可。

修正后的代码可以正确计算最小生成树的总权重。

全部评论

相关推荐

评论
点赞
收藏
分享

创作者周榜

更多
牛客网
牛客企业服务