研究了一个需要使用堆的题目,即和最小的k个数对
给定两个递增排序的整数数组,从两个数组中各取一个数字u和v组成一个数对(u, v),请找出和最小的k个数对。例如,输入两个数组[1, 5, 13, 21]和[2, 4, 9,15],和最小的3个数对为(1, 2)、(1, 4) 和 (2, 5)。
分析:观察一下解这套题逻辑,先把数组1的第一个数和数组2的第一个数组成一对,这就是第一个最小的数对。那么第二个呢?要么是数组1的第一个数和数组2的第二个数,即(1, 4);要么是数组1的第二个数和数组2的第一个数(5, 2)。通过计算,是(1, 4)。那么比(1, 4)大一点的数呢?那就是(1, 9)和(5, 4)……所以不断根据上一个最小的数对,把比上一个数大一点的数加入到队列中,再进行比较,直到找出所有的k个最小的数对。
注意,用这种逻辑会出现的问题是,数对中会有重复。
#include <iostream> #include <vector> #include <algorithm> #include <queue> using PII = std::pair<int*, int*>; struct greaterPtr { bool operator() (PII p1,PII p2) { if (*p1.first + *p1.second <= *p2.first + *p2.second) return 0; else return 1; } }; std::vector<PII> getKthPairOfNum(int *arr1, int m, int *arr2, int n, int k) { std::vector<PII> result; if (m <= 0 || n <= 0 || k <= 0) return result; if (arr1 == nullptr || arr2 == nullptr) return result; int *p1 = arr1; int *p2 = arr2; std::priority_queue<PII,std::vector<PII>, greaterPtr> minHeap; minHeap.push(std::pair<int*, int*>(p1,p2)); while (k != 0){ auto curPtr = minHeap.top(); p1 = curPtr.first; p2 = curPtr.second; minHeap.pop(); minHeap.push(std::pair<int*,int*>(p1+1, p2)); minHeap.push(std::pair<int*,int*>(p1, p2+1)); if (result.size() > 0){ auto last = result.back(); if (last.first == p1 && last.second == p2) continue; } result.push_back(curPtr); k--; } } void test1() { int arr1[] = {1, 5, 13, 21}; int arr2[] = {2, 4, 9, 15}; auto result = getKthPairOfNum(arr1, 4, arr2, 4, 3); int len = result.size(); for (auto i : result){ printf("(%d,%d)\n", *i.first, *i.second); } } int main() { test1(); }
这份代码中涉及到一个priority_queue在构造最小堆中,使用自定义的比较策略。需要传入一个函数对象作为第三个参数,可以是结构、类、或者模板类、模板函数、仿函数。