1. 程式人生 > >牛客網《劍指Offer》程式設計 29. 最小的k個元素 (使用堆)

牛客網《劍指Offer》程式設計 29. 最小的k個元素 (使用堆)

題目描述

輸入n個整數,找出其中最小的K個數。例如輸入4,5,1,6,2,7,3,8這8個數字,則最小的4個數字是1,2,3,4,。

解題思路

本題可以構建一個size為k的最大堆進行求解。

遍歷陣列。

當堆中元素總量小於k的時候,將當前元素放入堆;

當堆中元素總量等於k的時候,將當前元素與堆頂的元素比較。如果比最大的元素大,淘汰當前元素;如果比對頂元素小,對頂元素出堆,並將當前元素插入堆中。

該演算法的空間複雜度為O(k),時間複雜度為O(n)。

程式碼實現

class Solution {
public:
   void maxHeapInsert(vector<int> &maxHeap, int key) {//將插入的元素上調,相當於AdjustUp
	//前提是該陣列已經符合最大堆的要求了,因此插入一個數之後,要將該陣列調整成為最大堆。
	maxHeap.push_back(key);
	int idx = maxHeap.size()-1;
	for (; idx != 1 && key > maxHeap[idx / 2]; idx /= 2) {
		maxHeap[idx] = maxHeap[idx / 2];
	}
	maxHeap[idx] = key;
}
void maxHeapDelete(vector<int> &maxHeap) {//將堆頂的元素向下調整,最後調整成為最大堆。相當於AdjustDown
	if (maxHeap.size() == 1) {
		return;
	}
	if (maxHeap.size() == 2) {
		maxHeap.pop_back();
		return;
	}
	maxHeap[1] = maxHeap[maxHeap.size() - 1];
	maxHeap.pop_back();//用陣列中最後的一個元素覆蓋最大的元素,然後調整陣列使其符合最大堆的定義。
	size_t parent = 1, child = 2;
	int key = maxHeap[1];
	for (; child < maxHeap.size(); child *= 2)
	{
		if (child + 1 < maxHeap.size() && maxHeap[child] < maxHeap[child + 1])
		{
			child += 1;
		}
		if (maxHeap[child] < key) {
			break;
		}
		else {
			maxHeap[parent] = maxHeap[child];
			parent = child;
		}
	}
	maxHeap[parent] = key;
}
vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
    if(input.size()<k){
        return vector<int>();
    }
	vector<int> maxHeap;//maxHeap是從1開始計數的。
	maxHeap.push_back(0);
	for (int i = 0; i < input.size(); i++) {
		if (maxHeap.size() < k+1) {
			maxHeapInsert(maxHeap,input[i]);
		}
		else if(maxHeap.size()==k+1&&maxHeap[1]>input[i]){
			maxHeapDelete(maxHeap);
			maxHeapInsert(maxHeap, input[i]);
		}
	}
    //vector<int> res(maxHeap.begin()+1,maxHeap.end());
	return vector<int>(maxHeap.begin()+1,maxHeap.end());
}
};