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;