[各种面试题] 两个数组和的第K大

这是谷歌的一道面试题，有两个数组A和B，假设有一个数组C，C[i] = A[j] + B[ k ] ， 即C中的元素是A和B中两个元素的和。

让你求C中第K大的数字。

之前有一篇转载的用堆来求的方法，因为每出堆一次最多添加两个元素进来，所以堆的最大容量是 2* k， 所以入堆出堆复杂度是 logK, 最后的复杂度是KlogK。

注意添加的时候要判断是否已经进过堆。

#include<iostream>
#include<cstdio>
#include<vector>
#include<queue>
#include<string>
#include<cstring>
#include<climits>
#include<set>
#include<algorithm>
using namespace std;

struct node
{
	int a,b,val;
	node(int aa,int bb,int v):a(aa),b(bb),val(v){}
	bool operator>(const node& oth)const
	{
		return val>oth.val;
	}
};

int findKthSum(int A[],int m,int B[],int n,int k)
{
	priority_queue<node,vector<node>,greater<node> > Q;
	Q.push(node(0,0,A[0]+B[0]));
	set<pair<int,int> > visited;
	visited.insert(pair<int,int>(0,0));
	while(!Q.empty())
	{
		node t=Q.top();
		Q.pop();
		k--;
		if(k==0)
			return t.val;
		set<pair<int,int> >::iterator it;
		if(t.a+1<m&&(it=visited.find(pair<int,int>(t.a+1,t.b)))==visited.end())
		{
			visited.insert(it,pair<int,int>(t.a+1,t.b));
			Q.push(node(t.a+1,t.b,A[t.a+1]+B[t.b]));
		}
		if(t.b+1<n&&(it=visited.find(pair<int,int>(t.a,t.b+1)))==visited.end())
		{
			visited.insert(it,pair<int,int>(t.a,t.b+1));
			Q.push(node(t.a,t.b+1,A[t.a]+B[t.b+1]));
		}
	}
	return -1;
}
int main()
{
	int m,n,k;
	while(scanf("%d%d%d",&m,&n,&k)!=EOF)
	{
		int* A=new int[m];
		int* B=new int[n];
		for(int i=0;i<m;i++)
			scanf("%d",&A[i]);
		for(int i=0;i<n;i++)
			scanf("%d",&B[i]);
		sort(A,A+m);
		sort(B,B+n);

		int ans=findKthSum(A,m,B,n,k);
		printf("%d\n",ans);
	}
}

然后还有一种比较巧妙的方法，用二分答案然后判定来做。排序后，最小和是A[0]+B[0],最大是A[m-1]+B[n-1] ，所以在这个范围内二分答案，然后判定。

判定的准则是对于二分到的目标值 Piv, 在C里需要有至少K个和要小于它，如果满足，那么它就是可能的答案，但并不一定是，因为piv并不一定出现在C中，所以要记录它，然后继续二分。

计算小于的个数的时候，充分利用A和B有序，可以做到O（m+n），所以加上二分的复杂度，最后复杂度是  O ( logMAXSUM * (m+n ) ）

#include<iostream>
#include<cstdio>
#include<vector>
#include<string>
#include<cstring>
#include<climits>
#include<algorithm>
using namespace std;

long long  countSmaller(long long A[],long long m,long long B[],long long n,long long piv)
{
	long long pa=0,pb=n-1;
	long long cnt=0;
	for(;pa<m;pa++)
	{
		if(A[pa]>piv)
			break;
		while(pb>=0&&A[pa]+B[pb]>piv)
			pb--;
		cnt+=pb+1;
	}
	return cnt;
}
long long findKthSum(long long A[],long long m,long long B[],long long n,long long k)
{
	long long  l=A[0]+B[0];
	long long r=A[m-1]+B[n-1];
	long long ans=-1;
	while(l<=r)
	{
		long long mid=l+((r-l)>>1);
		if( countSmaller(A,m,B,n,mid)>=k )
		{
			ans=mid;
			r=mid-1;
		}
		else
			l=mid+1;
	}
	return ans;
}
int main()
{
	long long m,n,k;
	while(scanf("%lld%lld%lld",&m,&n,&k)!=EOF)
	{
		long long* A=new long long[m];
		long long* B=new long long[n];
		for(long long i=0;i<m;i++)
			scanf("%lld",&A[i]);
		for(long long i=0;i<n;i++)
			scanf("%lld",&B[i]);
		sort(A,A+m);
		sort(B,B+n);
		long long ans=findKthSum(A,m,B,n,k);
		printf("%lld\n",ans);
	}
}