从堆到TopK：一文吃透核心原理与实战应用

// 向上调整算法
void AdjustUp(size_t child_idx) {
	// 循环判断 维持堆的结构
	while (child_idx > 0) {
		// 找父节点索引
		size_t parent_idx = (child_idx - 1) / 2;

		if (_comp(_arr[parent_idx], _arr[child_idx])) {
			std::swap(_arr[parent_idx], _arr[child_idx]);
			child_idx = parent_idx;
		}
		else {
			break;
		}
	}
}

// 向下调整算法
void AdjustDown(size_t parent) {
	// 假设左孩子是最大的值 先找到左孩子索引
	size_t max_idx = 2 * parent + 1;

	// 循环判断 维持堆的正确结构
	while (max_idx < _size) {
		// 验证假设是否成立
		if (max_idx + 1 < _size && _comp(_arr[max_idx], _arr[max_idx + 1])) {
			++max_idx;// 更新索引 右孩子才是更大的值
		}

		// 现在已经找出两个孩子中最大的一个了 和父节点比较
		if (_comp(_arr[parent], _arr[max_idx])) {
			// 交换二者位置
			std::swap(_arr[max_idx], _arr[parent]);
			// 更新索引
			parent = max_idx; // 父节点索引下移到孩子节点
			max_idx = parent * 2 + 1; // 还是先假设左孩子的值更大
		}
		else {
			break;
		}
	}
}

// 构造函数：通过数组初始化堆（核心建堆操作）
max_heap(const std::vector<T>& arr) {
	_arr = arr;
	_size = arr.size();
	if (_size <= 1) return;  // 空堆或单个元素无需调整

	// 从最后一个非叶子节点开始，向前依次调整每个节点
	// 最后一个非叶子节点 = (最后一个节点索引 - 1) / 2
	for (int i = (_size - 2) / 2; i >= 0; --i) {
		AdjustDown(i);  // 对每个非叶子节点执行向下调整
	}
}

#include <iostream>
#include <cassert>
#include <vector>
#include <functional>

namespace Vect {
	template <class T,class Compare = std::less<T>>
	class heap {
	private:
		std::vector<T> _arr; // 底层是动态数组
		size_t _size; // 有效数据个数
		size_t _capacity; // 容量
		Compare _comp; // 决定大顶堆还是小顶堆


		// 向上调整算法
		void AdjustUp(size_t child_idx) {
			// 循环判断 维持堆的结构
			while (child_idx > 0) {
				// 找父节点索引
				size_t parent_idx = (child_idx - 1) / 2;

				if (_comp(_arr[parent_idx], _arr[child_idx])) {
					std::swap(_arr[parent_idx], _arr[child_idx]);
					child_idx = parent_idx;
				}
				else {
					break;
				}
			}
		}


		// 向下调整算法
		void AdjustDown(size_t parent) {
			// 假设左孩子是最大的值 先找到左孩子索引
			size_t max_idx = 2 * parent + 1;

			// 循环判断 维持堆的正确结构
			while (max_idx < _size) {
				// 验证假设是否成立
				if (max_idx + 1 < _size && _comp(_arr[max_idx], _arr[max_idx + 1])) {
					++max_idx;// 更新索引 右孩子才是更大的值
				}

				// 现在已经找出两个孩子中最大的一个了 和父节点比较
				if (_comp(_arr[parent], _arr[max_idx])) {
					// 交换二者位置
					std::swap(_arr[max_idx], _arr[parent]);
					// 更新索引
					parent = max_idx; // 父节点索引下移到孩子节点
					max_idx = parent * 2 + 1; // 还是先假设左孩子的值更大
				}
				else {
					break;
				}
			}
		}

		// 扩容
		void resize() {
			if (_size >= _capacity) {
				size_t new_cap = _capacity == 0 ? 4 : 2 * _capacity;
				_arr.resize(new_cap); // 直接调用vector的扩容接口 不用自己手动实现
				_capacity = new_cap;
			}
		}
	public:
		// 默认构造
		heap():_size(0),_capacity(0){}

		// 构造函数：通过数组初始化堆（核心建堆操作）
		heap(const std::vector<T>& arr) {
			_arr = arr;
			_size = arr.size();
			if (_size <= 1) return;  // 空堆或单个元素无需调整

			// 从最后一个非叶子节点开始，向前依次调整每个节点
			// 最后一个非叶子节点 = (最后一个节点索引 - 1) / 2
			for (int i = (_size - 2) / 2; i >= 0; --i) {
				AdjustDown(i);  // 对每个非叶子节点执行向下调整
			}
		}

		// 拷贝构造 vector自带深拷贝
		heap(const heap& other)
			:_size(other._size)
			,_capacity(other._capacity)
			,_arr(other._arr)
		{ }

		// 析构 vector自带
		~heap(){ }

		/*==================核心接口=========================*/

		// 插入元素（向上调整）
		void push(const T& val) {
			// 先扩容
			resize();         
			_arr[_size] = val; 

			// 向上调整新元素到正确位置
			AdjustUp(_size);   
			++_size;           
		}

		// 删除堆顶元素（向下调整）
		void pop() {
			// 确保不为空才能pop
			assert(!empty());

			// 交换堆顶和堆底元素 
			std::swap(_arr[0], _arr[_size - 1]);
			// 删除堆底元素
			--_size;

			// 向下调整
			AdjustDown(0);
		}

		// 获取堆顶元素
		const T& top()const { assert(!empty()); return _arr[0]; }

		// 判空
		const bool empty() const { return _size == 0; }

		// 获取元素个数
		const size_t size() const { return _size; }

		// 清空堆
		void clear() { _size = 0; }

		// 打印堆
		void PrintHeap() {
			if (empty()) {
				std::cout << "堆为空" << std::endl;
				return;
			}

			for (size_t i = 0; i < _size; i++)
			{
				std::cout << _arr[i] << " " ;
			}
			std::cout << "\n";
		}	
	};
    
    /*=================== 测试代码 =========================*/

void TestAPI() {
	std::cout << "===== 测试大根堆（int类型） =====" << std::endl;

	// 测试默认构造 + push + top
	heap<int> heap1;
	heap1.push(5);
	heap1.push(3);
	heap1.push(8);
	heap1.push(10);
	heap1.push(7);
	heap1.PrintHeap(); // 输出：堆的元素（数组存储）：10 8 5 3 7 （堆顶：10）

	// 测试pop（删除堆顶）
	heap1.pop();
	heap1.PrintHeap(); // 输出：堆的元素（数组存储）：8 7 5 3 （堆顶：8）

	// 测试string类型堆
	std::cout << "\n===== 测试string类型堆 =====" << std::endl;
	heap<std::string> heap3;
	heap3.push("apple");
	heap3.push("banana");
	heap3.push("cherry");
	heap3.push("date");
	heap3.PrintHeap(); // 输出：堆的元素（数组存储）：date cherry banana apple （堆顶：date）
	heap3.pop();
	std::cout << "Pop后堆顶：" << heap3.top() << std::endl; // 输出：cherry
 }
}

// TopK问题

// 找前K大 小顶堆维护
static std::vector<int> TopK_max(const std::vector<int>& data, size_t K) {
	// 边界处理 K=0或者K>=数据量直接返回该数组
	if (K == 0 || K >= data.size()) {
		return data;
	}

	// 小顶堆（greater 父 < 子）
	heap<T, std::greater<T>> minHeap;

	// 遍历数组
	for (size_t i = 0; i < data.size(); i++)
	{
		if (minHeap.size() < K) {
			// 堆中不足K个元素 直接放进去
			minHeap.push(data[i]);
		}
		else if (data[i] > minHeap.top()) {
			// 当前元素比堆顶大
			// 说明堆顶的最小值不够资格进前K 要淘汰
			minHeap.pop(); // 淘汰堆顶
			minHeap.push(data[i]); // 插入新元素
		}
		// 否则 data[i] <= 堆顶 直接丢弃 连第K大都排不上
	}

	// 此时堆中正好保存了无序的前K大元素
	std::vector<T> ret_nums;
	while (!minHeap.empty()) {
		ret_nums.push_back(minHeap.top());
		minHeap.pop();
	}
	return ret_nums;
}


// 找前K小 用大顶堆维护
static std::vector<T> TopK_min(const std::vector<T>& data, size_t K) {
	if (K == 0 || K >= data.size()) return data;

	// 大顶堆（less 父 > 子）
	heap<T, std::less<T>> maxHeap;

	// 遍历数组
	for (size_t i = 0; i < data.size(); i++)
	{
		if(maxHeap.size() < K){
			// 堆中不足K个元素 直接放
			maxHeap.push(data[i]);
		}
		else if (data[i] < maxHeap.top()) {
			// 如果当前元素比堆顶还小
			// 说明堆顶最大值不够资格进前K小 淘汰
			maxHeap.pop(); // 淘汰堆顶
			maxHeap.push(data[i]); // 插入新元素
		}
		// 否则 如果data[i] >= 堆顶 直接丢弃 连第K小都排不上
	}

	// 此时堆中正好保存了无序的前K小元素
	std::vector<T> ret_nums;
	while (!maxHeap.empty()) {
		ret_nums.push_back(maxHeap.top());
		maxHeap.pop();
	}
	return ret_nums;
}

void TestTopK() {
	std::vector<int> nums = { 5, 3, 8, 10, 7, 2, 9 ,100,2058,60,102,5461};

	auto top3max = heap<int>::TopK_max(nums, 3);
	std::cout << "前3大：";
	for (auto e : top3max) std::cout << e << " ";
	std::cout << "\n";

	auto top3min = heap<int>::TopK_min(nums, 3);
	std::cout << "前3小：";
	for (auto e : top3min) std::cout << e << " ";
	std::cout << "\n";
}