华为OJ2051-最小的K个数（Top K问题）

发表于   |   分类于   |  

一、题目描述

描述：

输入n个整数，输出其中最小的k个。

输入：

  1. 输入 n 和 k
  2. 输入一个整数数组

输出：

输出一个整数数组

样例输入：

1
2
5 2
1 3 5 7 2

样例输出：

1
1 2


二、Top K问题

对于 Top K 问题有很多种解法。

解法一：排序

相信很多人会首先想到这种方法，先把数组按升序/降序进行排序，然后输出 K 个最小/最大的数。

  • 常规的排序方法时间复杂度至少是$Θ(nlog_2n)$。（快排或堆排序）
  • 可能你会说，我们可以使用线性时间的排序算法。当然可以，但通常它们对输入的数组有一定的要求。比如计数排序要求 n 个数都是正整数，且它们的取值范围不太大。

解法二：部分排序 $O(n*k)$

由于我们只需要找出最小/最大的 k 个数，所以我们可以进行部分排序，比如简单选择排序 和 冒泡排序，它们每一趟都能把一个最小/最大元素放在最终位置上，所以进行 k 趟就能把 n 个数中的前 k 个排序出来。

部分简单选择排序：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
void select_sort(int A[], int n, int k)
{
	for(int i=0; i<k; ++i) { // k趟
		int Min = i;         // 记录最小元素的位置
		
		for(int j=i+1; j<n; ++j)
			if(A[j] < A[Min])
				Min = j;

		if(Min != i)  // 与A[i]交换
		{
			int tmp = A[Min];
			A[Min] = A[i];
			A[i] = tmp;
		}
	}
}

部分冒泡排序：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
void bubble_sort(int A[], int n, int k)
{
	for(int i=0; i<k; ++i)  // k趟
	{
		bool flag = false;
		for(int j=n-1; j>i; --j)  // 一趟冒泡过程
			if(A[j-1] > A[j])
			{
				int tmp = A[j-1];
				A[j-1] = A[j];
				A[j] = tmp;
				flag = true;
			}
		if(flag == false)  // 已经有序
			return ;
	}
}

那么，$O(nlog_2n)$ 与 $O(n*k)$ 哪一个更好呢？这取决于 k 的大小。在 k 较小的情况下，即 $k <= log_2n$，可以选择部分排序。

解法三：快排划分 $O(n*log_2k)$

根据基于快排partition操作的《第k顺序统计量的求解》，我们知道，当我们求出第 k 顺序统计量时，位于它前面的元素都比它小，位于它后面的元素都比它大。这时，数组的前 k 个数就是最小的 k 个数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
int partition(int A[], int low, int high)
{
	int pivot = A[low];
	while(low < high)
	{
		while(low < high && A[high]>=pivot)
			--high;
		A[low] = A[high];
		while(low < high && A[low]<=pivot)
			++low;
		A[high] = A[low];
	}
	A[low] = pivot;
	return low;
}


int topK(int A[], int low, int high, int k)
{
	if(k <= 0)
		return -1;
	if(low == high)
		return low;

	int pos = partition(A, low, high);
	int i = pos - low + 1;
	if(i == k)
		return pos;  // 返回前k个数的
	else if(i > k)
		return topK(A, low, pos, k);
	else
		return topK(A, pos+1, high, k-i);
}

我们说这个算法的平均时间复杂度是线性的，更准确地说，是 $O(n∗log_2k)$。另外，为了避免特殊数据下的算法退化，最好使用随机化版本的划分操作。

解法四：大根堆 $O(n∗log_2k)$

参见《堆排序》，可以用大小为 k 的大根堆来存储最小的 k 个数。大根堆的堆顶元素就是最小 k 个数中最大的一个。每次新考虑一个数 X：

  • 如果 X 比堆顶的元素 Y 大，则不需要改变原来的堆，因为这个元素比最小的 k 个数都大。

  • 如果 X 比堆顶元素 Y 小，那么用 X 替换堆顶的元素 Y。在 X 替换堆顶元素 Y 之后，大根堆的结构可能被破坏，需要进行向下调整。调整过程的时间复杂度是 $O(log_2k)$ 。

遍历完成以后，数组的前 k 个数就是最小的 k 个数，但是它们并非有序，而是以堆的形式存在。C++代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
void AdjustDown(int A[], int i, int len)  
{  
	int temp = A[i];  // 暂存A[i]  

	for(int largest=2*i+1; largest<len; largest=2*largest+1)  
	{  
		if(largest!=len-1 && A[largest+1]>A[largest])  
			++largest;         // 如果右子结点大  
		if(temp < A[largest])  
		{  
			A[i] = A[largest];  
			i = largest;         // 记录交换后的位置  
		}  
		else  
			break;  
	}  
	A[i] = temp;    // 被筛选结点的值放入最终位置  
}

/* 建堆 */
void BuildMaxHeap(int A[], int len)
{
	for(int i=len/2-1; i>=0; --i)  // 从i=n/2-1到0，反复调整堆
		AdjustDown(A, i, len);
}


/* 维护 A[0...k-1] 这个大根堆 */
void topK(int A[], int n, int k)
{
	BuildMaxHeap(A, k);  // 先用前面的k个数建大根堆
	for(int i=k; i<n; ++i)
	{
		if(A[i] < A[0])  // 如果小于堆顶元素，替换之
		{
			int tmp = A[0];
			A[0] = A[i];
			A[i] = tmp;
			AdjustDown(A, 0, k);  // 向下调整
		}
	}
}

注意：找最小的 $k$ 个数，就维护一个大根堆；找最大的 $k$ 个数，就维护一个小根堆。


三、解题报告

第二部分已经讲解地很清楚了，几种解法都可以，只要注意输入输出的格式就行了。