Description:
两个长度为 n 的数组 A 和 B, 各从中选出一个元素相加 A[i] + B[j], 求 top n 小的那些和。
思路 1:这样的和总共有 n^2 个, 排序,然后取前 n 个。
时间复杂度 O(n^2 lgn), 空间复杂度 O(n^2)
思路 2:用 最大堆。
1. 首先用 n 个和建立一个最大堆,
2. 然后循环剩余的 n^2 - n 个和,依次与堆顶元素比较(最大值), 如果小于最大值,则弹出堆顶,push 进入当前的和。
时间复杂度 O(n^2 lgn), 空间复杂度是 O(n)
思路3: 还是采用 最小堆,并采用 多路合并 的思路。
首先,分析多路合并问题: 把 K 个有序数组合并成一个数组。
1. 将 K 个数组的首元素(最小),建立一个最小堆;
2. 弹出最小值,并插入新的数组中;将最小值所在的数组的下一个元素插入堆中。这样依次执行下去。
多路合并的代码见 http://blog.csdn.net/shoulinjun/article/details/19576585
回到我们的问题,将 N^2 个 sum 写成如下的形式:
所有 N^2 个和:(数组 A 和 B 排序后)
A[0] + B[0] <= A[0] + B[1] <= A[0] + B[2] <= ... <= A[0] + B[n-1]
A[1] + B[0] <= A[1] + B[1] <= A[1] + B[2] <= ... <= A[1] + B[n-1]
...............................
A[n-1] + B[0] <= A[n-1] + B[1] <= A[n-1] + B[2] <= ... <= A[n-1] + B[n-1]
注意: 第一步将数组 A 和 B 排序,则上述的 N 个序列都是有序的。
Bingo! 这时就可以采用与 多路合并 同样的算法:
Algorithm:
1. 将数组 A 和 B 排序;
2. 用 上述 N 个序列的首元素(最小),建立一个最小堆;
3. 弹出最小值,即为 top N;将最小值所在的数组的下一个元素插入堆中。这样执行 N 次, 得到 所有的 top N。
Complexity:
O(NlgN) + O(N)+ N*O(lgN) = O(NlgN)
Implementation:
实现细节的话,算法的第 2 步中, 最小值所在的数组的下一个元素如何得到?? 可以同时记录最小值这个和中所包含的B 中元素的下标。
下一个元素的值就是 min - B[i] + B[i+1]
//copyright @ L.J.SHOU Feb.14, 2014
#include <iostream>
#include <vector>
#include <utility>
#include <queue>
#include <algorithm>
#include <iterator>
using namespace std; //A and B have the same size K
vector<int> TopKSum(vector<int> &A, vector<int> &B)
{
vector<int> result;
priority_queue<pair<int, int>, vector<pair<int, int> >, greater<pair<int, int> > > Q; /* sort */
sort(A.begin(), A.end());
sort(B.begin(), B.end()); for(int i=0; i<A.size(); ++i)
Q.push(make_pair(A[i] + B[0], 0)); /* 0 denotes index in B */ for(int i=0; i<A.size(); ++i)
{
pair<int, int> e = Q.top(); Q.pop();
result.push_back(e.first);
/* push next element back into heap */
e.first = e.first - B[e.second] + B[e.second+1];
e.second ++;
Q.push(e);
} return result;
} int main(void)
{
vector<int> A, B, sum;
A.push_back(1);
A.push_back(2);
A.push_back(3); B.push_back(2);
B.push_back(4);
B.push_back(4); sum = TopKSum(A, B);
/* print vector */
copy(sum.begin(), sum.end(), ostream_iterator<int>(cout, " "));
cout << endl; return 0;
}
思考题:
n 个长度为 n 的数组,每个数组中选出一个元素相加, 求 top n 小的那些 sum。
复杂度: n^2 lgn