TopK 问题——思路总结

面试常考题 $n$ 个数，取前 $k$ 小的数：
总结来说

• 当 $n$ 非常大时，且动态，而 K固定时，【下面两个方法时间复杂度都是 $O(n\log (K))$】
  • 我们可以使用 二叉堆 的方法（容量为 $K$ 的大顶堆），比Treemap快，但缺点K个数是没有顺序
  • 还可以使用 Treemap 的方法（容量为 $K$ 的红黑树），K 个数排好了序
• 当 $n$ 非常大，K 也比较大，我们使用快排剪枝的方法 【平均时间复杂度 $O(n)$，空间$O(\log (n)$】
• 当 $n$ 非常大，但是数组里面的元素取值范围给出来了，例如$arr[i]\in (0,100)$，这时候直接 计数排序，遍历一遍即可得到答案 【时间复杂度：$O(n)$】

剑指 Offer 40. 最小的k个数

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

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]
示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]
 

解法一：直接排序

vector<int> getLeastNumbers(vector<int>& arr, int k) {
	vector<int> ans;
	sort(arr.begin(), arr.end());
	ans.assign(arr.begin(), arr.begin() + k);
	return ans;
}

• 时间复杂度：$O(n\log n)$
• 空间复杂度：$O(\log n)$

解法二：Treemap统计频次

Key 是数字，value 是该数字出现的次数；
我们遍历数组中的数字，维护一个数字总个数为 K 的 TreeMap：

1. 若目前 map 中数字个数小于 K，则将 map 中当前数字对应的个数 +1；
2. 否则，判断当前数字与 map 中最大的数字的大小关系：
   2.1 若当前数字大于等于 map 中的最大数字，就直接跳过该数字；
   2.2 若当前数字小于 map 中的最大数字，则将 map 中当前数字对应的个数 +1，并将 map 中最大数字对应的个数减 1。

注：map::iterater 在遍历的时候是按照 $key$ 值的升序顺序迭代的，相当于中序遍历了平衡树。从上面代码可以看到如何在O(1)的时间内完成操作。如果需要找最大的 $key$ , 则只需要如下代码

auto it = map.end();
it--;
cout << "最大的key值: " << it->first << endl;

• 时间复杂度：$O(n\log (K))$，红黑树的每次操作（生成树、查找树）都是$\log (K)$ 的复杂度
• 空间复杂度：$O(K)$， K 个节点的红黑树

解法三：大顶堆（求前k个小的元素）

vector<int> getLeastNumbers(vector<int>& arr, int k) {
	vector<int> ans;
    if (k==0) return ans;
	priority_queue<int, vector<int>, less<int>> q;
	for (int i = 0; i < arr.size(); i++) {
		if (q.size() < k) q.push(arr[i]);
		else {
			if (q.top() > arr[i]) {
				q.pop(); q.push(arr[i]);
			}
		}
	}
	while (!q.empty()) {
		ans.push_back(q.top()); q.pop();
	}
	return ans;

• 时间复杂度：$O(n\cdot log(k))=O(n)$，因为我们只需要生成 $k$ 个节点的二叉堆
• 空间复杂度：$O(k)$

解法四：基于快排的剪枝优化

快排思路：

1. 用一个partition作为分割符，分出两个数组，前一个数组 A 中的元素全部小于 partition 所在位置的元素，后一个数组 B 中的元素全部大于 partition 所在位置的元素。
2. 然后分别递归这两个数组。

当我们下一轮的partition的位置>k时，说明我们需要的前 k 个元素存在 A 中，此时我们剪枝 B 的递归
当我们下一轮的partition的位置<k时，说明我们需要的前 k 个元素存在 B 中，此时我们剪枝 A 的递归

void qucikSortPlus(vector<int>& arr, int l, int r, int k) {
	if (l >= r) return;
	int left = l, right = r;
	int partition = l;
	while (l < r) {
		while (arr[r] >= arr[partition] && l < r) r--;
		while (arr[l] <= arr[partition] && l < r) l++;
		swap(arr[l], arr[r]);
	}
	swap(arr[partition], arr[l]);
	// 剪枝，
	if (l < k) qucikSortPlus(arr, l + 1, right, k);
	if (l > k) qucikSortPlus(arr, left, l - 1, k);
}
vector<int> getLeastNumbers(vector<int>& arr, int k) {
	vector<int> ans(k);
	qucikSortPlus(arr, 0, arr.size() - 1, k);
	ans.assign(arr.begin(), arr.begin() + k);
	return ans;
}

快排优化：随机选取partition

   int partition(vector<int>& nums, int l, int r) {
        int pivot = nums[r];
        int i = l - 1;
        for (int j = l; j <= r - 1; ++j) {
            if (nums[j] <= pivot) {
                i = i + 1;
                swap(nums[i], nums[j]);
            }
        }
        swap(nums[i + 1], nums[r]);
        return i + 1;
    }

    // 基于随机的划分
    int randomized_partition(vector<int>& nums, int l, int r) {
        int i = rand() % (r - l + 1) + l;
        swap(nums[r], nums[i]);
        return partition(nums, l, r);
    }

    void randomized_selected(vector<int>& arr, int l, int r, int k) {
        if (l >= r) {
            return;
        }
        int pos = randomized_partition(arr, l, r);
        int num = pos - l + 1;
        if (k == num) {
            return;
        } else if (k < num) {
            randomized_selected(arr, l, pos - 1, k);
        } else {
            randomized_selected(arr, pos + 1, r, k - num);
        }
    }
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        srand((unsigned)time(NULL));
        randomized_selected(arr, 0, (int)arr.size() - 1, k);
        vector<int> vec;
        for (int i = 0; i < k; ++i) {
            vec.push_back(arr[i]);
        }
        return vec;
    }

• 时间复杂度: $O(2n)=O(n)$, 最坏情况下的时间复杂度为 $O(n^2)$。情况最差时，每次的划分点都是最大值或最小值，一共需要划分 $n-1$ 次，而一次划分需要线性的时间复杂度，所以最坏情况下时间复杂度为 $O(n^2)$
• 空间复杂度：期望为 $O(\log n)$ 递归调用的期望深度为 $O(\log n)$，每层需要的空间为 $O(1)$，只有常数个变量。最坏情况下的空间复杂度为 $O(n)$。最坏情况下需要划分 $n$ 次，即 randomized_selected 函数递归调用最深 $n-1$ 层，而每层由于需要 $O(1)$ 的空间，所以一共需要 $O(n)$ 的空间复杂度。

情景2：arr.size()非常大时，但是我们知道arr中的元素取值范围在（0，10000）之间

解法五：计数排序

该方法适合数据量巨大，但分布较为集中的情形
当我们已知arr[i]中的元素取值范围时，0 <= arr[i] <= 10000，我们可以直接用一个数组mem去统计元素出现的频次。直到满足 k个数即可。

vector<int> getLeastNumbers(vector<int>& arr, int k) {
	if (k == 0) return {};
	vector<int> ans(k);
	int mem[10001] = {};
	for (int i = 0; i < arr.size(); i++) {
		mem[arr[i]]++;
	}
	int idx = 0;
	for (int i = 0; i < 10001; i++) {
		while (mem[i]-- > 0 && k > idx) ans[idx++] = i;
		if (k == idx) break;
	}
	return ans;
}

• 时间复杂度：$O(n+K)$
• 空间复杂度：$O(K)$