算法（六）二叉堆获取最小的k个数

关键词：heap

如果你有一个文件，里面包含20万行的整数，如何获取前k个最小的数？首先可以想到两个思路：

1. 将所有的数按从小到大排序，取前k个。
2. 先读入前k个数到一个数组中（大小为k）并按从小到大排序，然后每读入一个新的数就将其放入数组中合适的排序位置。当所有的数都按这个规则被处理后，最终留在数组中的k个数就是我们想要的。

最近我学习了一种新的数据结构，那就是二叉堆（以下简称堆），用它来解决上述问题在时间上往往更高效。（具体代码见下文）

假设我们的文件 20w_int.txt 包含20万行整数，三种方法取前k个最小数用时如下：
（其中sort代表第一种思路，k-array代表第二种思路，heap代表堆这种思路）

可以看出sort和heap两种方法随着k增大其耗时变化不大，但是k-array这种方法随k增大其耗时快速增加。

第一种思路——sort
将所有的数按从小到大排序，取前k个。
直接用GNU sort就行，假设取前10个最小数：

sort -n 20w_int.txt | head -10

第二种思路——k-array
先读入前k个数到一个数组中（大小为k）并按从小到大排序，然后每读入一个新的数就将其放入数组中合适的排序位置。当所有的数都按这个规则被处理后，最终留在数组中的k个数就是我们想要的。
C代码如下：

#include <stdio.h>
#include <stdlib.h>
#define MAXCHAR 20

int _getline(FILE* fp, char* line);
void insertSort(int* a, const int n, int d);

int main(int argc, char* argv[]) {
	FILE* fp;	// data file
	int k;		// top k
	int* a;
	char s[MAXCHAR];
	int nline = 0;
	int i;
	if (argc < 3) {
		printf("Usage: %s <data> <k>\n", argv[0]);
		return 1;
	}
	if ((fp = fopen(argv[1], "r")) == NULL) {
		printf("Error: cannot open %s\n", argv[1]);
		return 2;
	}
	k = atoi(argv[2]);
	if ((a = (int*) malloc(sizeof(int) * (k + 1))) == NULL) {
		printf("Error: out of space!\n");
		return 3;
	}
	while (nline < k && _getline(fp, s) >= 0) {
		insertSort(a, nline, atoi(s));
		nline++;
	}
	if (nline < k) {
		printf("Error: %s less than %d lines\n", argv[1], argv[2]);
		return 4;
	}
	while (_getline(fp, s) >= 0)
		insertSort(a, k, atoi(s));
	for (i = 0; i < k; i++)
		printf("%d\n", a[i]);
	fclose(fp);
	free(a);
	return 0;
}

int _getline(FILE* fp, char* line) {
	int n, c;
	if (feof(fp)) return -1;
	for (n = 0; (c = fgetc(fp)) != '\n' && c != EOF; line[n++] = c)
		;
	line[n] = '\0';
	return n;
}

void insertSort(int* a, const int n, int d) {
// a: a sorted int array from min to max.
// n: current size of a.
// d: num to be inserted.
	int i, j;
	for (i = 0; i < n && a[i] <= d; i++)
		;
	for (j = n - 1; j >= i; j--)
		a[j + 1] = a[j];
	a[i] = d;
}

第三种思路——heap
二叉树的特点就是其父节点的值总是不大于其子节点的值。一次插入花费O(logN)时间。具体代码如下：

#include <stdio.h>
#include <stdlib.h>
#define MINNUM -99999999
#define MAXCHAR 20
#define MAXLINE 10000000

int _getline(FILE* fp, char* line);
struct Heap;
typedef struct Heap *heap;
heap Init(const int n);
void Insert(const int n, heap h);
int DeleteMin(heap h);
void Destroy(heap h);
int isEmpty(heap h);
int isFull(heap h);

struct Heap {
	int capacity;
	int size;
	int* arr;
};

int main(int argc, char* argv[]) {
	FILE* fp;
	heap h;
	char line[MAXCHAR];
	int i, k;
	if (argc < 3) {
		fprintf(stderr, "Usage: %s <file> <k>\n", argv[0]);
		exit(11);
	}
	if ((fp = fopen(argv[1], "r")) == NULL) {
		fprintf(stderr, "Error: cannot open %s\n", argv[1]);
		exit(12);
	}
	k = atoi(argv[2]);
	h = Init(MAXLINE);
	while (_getline(fp, line) >= 0)
		Insert(atoi(line), h);
	for (i = 0; i < k; i++)
		printf("%d\n", DeleteMin(h));
	fclose(fp);
	Destroy(h);
}

int _getline(FILE* fp, char* line) {
	int n, c;
	if (feof(fp)) return -1;
	for (n = 0; (c = fgetc(fp)) != '\n' && c != EOF; line[n++] = c)
		;
	line[n] = '\0';
	return n;
}

heap Init(const int n) {
	heap h;
	if ((h = (heap) malloc(sizeof(struct Heap))) == NULL) {
		fputs("Error: out of space!\n", stderr);
		exit(1);
	}
	if ((h->arr = (int*) malloc(sizeof(int) * (n + 1))) == NULL) {
		fputs("Error: out of space!\n", stderr);
		exit(2);
	}
	h->arr[0] = MINNUM;
	h->capacity = n;
	h->size = 0;
	return h;
}

void Insert(const int n, heap h) {
	int i, j;
	if (isFull(h)) {
		fputs("Error: heap is full.\n", stderr);
		exit(3);
	}
	for (i = ++h->size, j = i / 2; n < h->arr[j]; i = j, j = i / 2)
		h->arr[i] = h->arr[j];
	h->arr[i] = n;
}

int DeleteMin(heap h) {
	int i, j, k;
	if (isEmpty(h)) {
		fputs("Error: heap is empty.\n", stderr);
		exit(4);
	}
	k = h->arr[1];
	for (i = 1, j = 2; j < h->size; i = j, j *= 2) {
		if (h->arr[j] > h->arr[j + 1])
			j++;
		if (h->arr[j] >= h->arr[h->size]) {
			h->arr[i] = h->arr[h->size];
			break;
		} else {
			h->arr[i] = h->arr[j];
		}
	}
	if (j >= h->size)
		h->arr[i] = h->arr[h->size];
	h->size--;
	return k;
}

void Destroy(heap h) {
	free(h->arr);
	free(h);
}

int isEmpty(heap h) {
	return h->size <= 0 ? 1 : 0;
}

int isFull(heap h) {
	return h->size >= h->capacity ? 1 : 0;
}

（公众号：生信了）