手写堆&堆排序&TopK问题
TopK问题:最大K个用最小堆 -> 堆顶最小, 若比堆顶还小, 则可直接忽略 O(nlogk)
排序问题:升序用最大堆 -> 使最大值在堆顶, 然后置尾, 依次循环所有元素 O(nlogn)
关于自定义比较:比较函数同排序规则, 在adjust和TopK的堆顶比较中采用(前, 后)与(新, 顶)的方式
Leetcode相关:
1337. 矩阵中战斗力最弱的 K 行
347. 前 K 个高频元素
堆排序:升序 -> 最大堆, 每次将最大值置于堆顶, 再将最大值置尾, 依次迭代
class HeapSort { void heapSort(vector<int> &nums) { int n = nums.size(); for (int i = n / 2 - 1; i >= 0; i--) { // 从最后一个非叶子结点开始, 向上调整 adjust(nums, i, n); } for (int j = n - 1; j > 0; j--) { // 每次取最大堆头与尾巴交换, 再重新调整 swap(nums[j], nums[0]); adjust(nums, 0, j); } } void adjust(vector<int> &nums, int i, int n) { for (int j = i * 2 + 1; j < n; j = j * 2 + 1) { if (j + 1 < n && nums[j] < nums[j + 1]) { // 右儿子更大, 指向右儿子 j++; } if (nums[i] < nums[j]) { // 儿子更大, 父子交换, 继续迭代 swap(nums[i], nums[j]); i = j; } else { break; } } } };
TopK问题:最大K个 -> 最小堆, 若新元素大于堆顶, 则替换堆顶后调整
class TopK { public: vector<int> getTopK(vector<int> &nums, int k) { if (k <= nums.size()) { return nums; } // 前k个元素直接入堆, 然后从最后一个非叶子结点向上调整 for (int i = 0; i < k; i++) { minHeap.emplace_back(nums[i]); } for (int i = k / 2 - 1; i >= 0; i--) { adjust(i); } // 如果当前元素比堆顶要大, 则替换堆顶后调整堆顶 for (int i = k; i < nums.size(); i++) { if (nums[i] > minHeap.front()) { minHeap.front() = nums[i]; adjust(0); } } vector<int> res; while (!minHeap.empty()) { // 每次取出堆顶元素, 然后首尾互换, 弹出尾巴, 调整堆顶 res.emplace_back(minHeap.front()); swap(minHeap.front(), minHeap.back()); minHeap.pop_back(); adjust(0); } reverse(res.begin(), res.end()); // 最后是TopK升序, 需要逆置 return res; } void adjust(int i) { int n = minHeap.size(); for (int j = i * 2 + 1; j < n; j = j * 2 + 1) { if (j + 1 < n && minHeap[j] > minHeap[j + 1]) { // 右儿子更小, 指向右儿子 j++; } if (minHeap[i] > minHeap[j]) { // 儿子更小, 父子交换, 继续迭代 swap(minHeap[i], minHeap[j]); i = j; } else { break; } } } private: // 最大k个, 用最小堆 vector<int> minHeap; };