C++中priority_queue的简单用法

学习算法过程中经常会遇到堆，而STL中的priority_queue是（优先队列）就是一个封装好的堆结构。

题目描述

输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4,。

class Solution {
public:
	vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
		int size = input.size();
		vector<int> ret;
        
        if (size == 0 || k <= 0 || k > size)
            return ret;
        
		priority_queue<int> q;
		for (int i = 0; i < size; i++) {
			if (q.size() < k)
				q.push(input[i]);
			else {
				if (input[i] < q.top()) {
					q.pop();
					q.push(input[i]);
				}
			}
		}

		while (!q.empty()) {
			ret.push_back(q.top());
			q.pop();
		}

		reverse(ret.begin(), ret.end());

		return ret;
	}
};

priority_queue默认为大顶堆，即堆顶元素为堆中最大元素。如果我们想要用小顶堆的话需要增加使用两个参数：

priority_queue<int, vector<int>, greater<int> > q;  // 小顶堆
priority_queue<int, vector<int>, less<int> > q;     // 大顶堆

尽管提供了大小顶堆的模板，但是我们在实际应用过程中往往会使用更加复杂的数据结构构建堆，这时候我们就需要自定义数据结构的比较方式了。

方式一：

struct Node{
    int x, y;
    Node(int a = 0, int b= 0):x(a), y(b) {}
};
 
struct cmp{
    bool operator() (const Node& a, const Node& b ){
	    if (a.x == b.x) 
			return a.y > b.y;
	    return a.x > b.x; 
	}
};

priority_queue<Node, vector<Node>, cmp> q; 

方式二：

struct Node{
    int x, y;
    Node(int a = 0, int b= 0):x(a), y(b) {}
};

bool operator < (const Node& a, const Node& b ){
    if (a.x == b.x) 
		return a.y > b.y;
    return a.x > b.x; 
}
 
priority_queue<Node> q;