求两个有序整型数组元素和的第K大值

妖狐艹你老母 2022-08-07 17:45 189阅读 0赞

今天与人讨论问题，发现一个很有意思的题目：**给定两个有序的整型数组，要求在最优的情况下找到两个数组元素求和后的第K大的值。**

1、刚开始遇到这个问题，我直观的想法就是分配一个数组保存两个数组元素的和，然后利用找有序数组中第K大元素的方法进行求解。然后分析了一下时间和空间复杂度，假设原始的两个有序数组分别为A和B，对应的长度为m和n，那么A和B中元素和就有m\*n个，即我们需要分配一个新的数组C\[m\*n\]保存对应的元素和，填满C数组的时间开销为O(m\*n)，然后在C中找第K大的元素，最快的时间复杂度为线性的，所以该方案的时间复杂度为O(m\*n)，空间复杂度为O(m\*n)。

如果m和n很大时，时间复杂度与空间复杂度都会很大，那么我们能不能进行优化呢？一般我们求解一个数组中的元素都会采用“二分”的方法，那么我们这个题目能不能采用这种策略呢？

2、答案是肯定的。因为A和B是两个有序数组，所以**A\[0\]+B\[0\]会是所有和的最小值，**不妨记为**min**，同时，**A\[m-1\]+B\[n-1\]对应所有和的最大值**，不妨记为**max**，那么我们要求的第K大值肯定在\[min, max\]区间中，我们不妨**设第K大的值为C\[K\]**，所以我们可以在区间\[min, max\]通过二分查找C\[K\]，因为\[min, max\]之间可能有很多重复元素，所以我们在查找C\[K\]的过程中可以通过判断区间中有多少元素大于这个C\[K\]。**如果比C\[K\]大的元素小于K个，而且比其小的元素也不到m\*n-K个，那么C\[K\]就是我们要求的值。如果比C\[K\]大的元素超过K个，那么我们就只需要在\[mid, max\]中找第K大的元素。如果比C\[K\]大的元素小于K个，那么需要在\[min, mid-1\]中找第(K-Count(C\[K\] + mid))大的元素**。这样原问题就转化为统计多少个数的和大于C\[K\]。

对于这个问题，可以采取遍历其中一个数组的方式，假如遍历数组A，然后利用二分在B中找C\[K\]-A\[i\]对应的元素，从而计算出相应的个数。因为A，B有序，所以如果找到A\[i\]+B\[j\]>C\[K\]，那么A\[p\] + B\[j\] (i<p<=m)都会大于C\[K\]，同理，A\[i\] + B\[q\] (j<q<=n)都会大于C\[K\]。利用这个规则可以简化我们的计数。

程序JAVA源代码如下：

<span style="white-space:pre">	</span>public int getKthSum(int[] A, int[] B, int K) {
    		if(A.length == 0 && B.length == 0)
    			return 0; 
    		if(A.length == 0) 
    			return K>1 ? B[K-1] : 0;
    		if(B.length == 0)
    			return K>1 ? A[K-1] : 0;
    		if(K == 0)
    			return 0;
    		int m = A.length - 1;
    		int n = B.length - 1;
    		int min = A[0] + B[0], max = A[m] + B[n];
    		while(min < max) {
    			int mid = (max + min) >>> 1;
    			int count = Count(mid, A, m, B, n);
    			if(count >= K) {
    				min = mid + 1;
    			} else {
    				max = mid;
    			}
    		}
    		return min;
    	}
    	
    	/**
    	 * 计算A和B的和中比value大的元素的个数有多少
    	 * @param value
    	 * @return
    	 */
    	public int Count(int value, int[] A, int m, int[] B, int n) {
    		int count = 0;
    		for(int i=0, j=n; i<=m && j>=0;) {
    			if(A[i] + B[j] > value) {
    				count += m - i + 1;  //因为A[p]+B[j]>value, i<p<=m
    				j--;
    			} else 
    				i++;
    		}
    		return count;
    	}

时间复杂度分析：二分法时间复杂度为O(log(max-min))，计算比C\[K\]大的元素个数的时间开销为O(m+n)，因为两者为嵌套关系，所以总的 **时间复杂度为O((m+n)log(max-min))**。

**空间复杂度为O(1)**。