一道Google面试题 ---- 求矩阵中第K小元素

近日另一友人也去google面试，被问及一算法题如下：

x，y为两个有序数组（升序），长度分别为M和N，矩阵m为一M*N的两维数组，其中m[i][j] = x[i] + y[j]。 设计一算法求矩阵中第K小的元素。

我所想的算法如下，与大家分享交流，看看有什么可以改进的。首先这题让我直觉上就想到了merge sort。其实每个m[i] （ 0<=i<M ）本身就是一个升序的一维数组。也就是说m是由M个有序数组组成，只要两两对这两个有序数组做merge sort，取前K个结果即可，代码的时间复杂度应该是O(K * min(M,N) )

 

import java.util.Arrays;

public class KthInMatrix {

	
	//初始化两个输入的升序数组
	static int x[] = {2,4,9,10,23,77};
	static int y[] = {6,9,11,22,34,45,78};
	static int K = 10;
	static int M = x.length;
	static int N = y.length;
	static int MIN = 0;
	static int MAX = 0;
	/**
	 * @param args
	 */
	public static void main(String[] args) {
		// TODO Auto-generated method stub
		System.out.println(getKthInMatrix());
		
		
	}
	//找到第K个小的数，如果K大于M*N，返回Integer.MAX_VALUE
	public static int getKthInMatrix()
	{
		
		printMatrix();
		
		//决定换行merge还是按列merge,按小的那个merge,这里分两个方法做，以减少判断带来的开销
		//虽然带码会有些冗余，但应该performance会好点
		
		if ( M > N)
		{
			MIN = N;
			MAX = M;
			int[] temp = x;
			x = y;
			y = temp;
			
		} 
		else
		{
			
			MIN = M;
			MAX = N;
			
			
		}
		return mergebyMin();
	}
	
	public static int mergebyMin()
	{
		//存放K个最小的数
		int temp[] = new int[K];
		int result[] = new int[K];
		Arrays.fill( result , Integer.MAX_VALUE);
		
		
	    
		for (int i = 0  ; i < MIN ; i ++ )
		{
			
			int j = 0, k = 0, index = 0;
			
			//j, k , index 分别为游历 m[i] , temp 和 result的坐标 , index < K 保证了 k < K
			while (j < MAX && index < K)
			{
				
				//计算当前坐标的值
				int m = x[i] + y[j];
				
				if ( m < result[k])
				{
					temp[index ++] = m;
					j ++ ; 
					
				} else
				{
					temp[index ++] = result[k ++];
				}
				
				
			}
			
			while ( index < K )
			{
			if (j < MAX)
			    temp[index ++] = x[i] + y[j ++];
			else
				temp[index ++] = result[k ++];
			}
			result = temp;
			temp = new int[K];
			
		}
		
		return result[K-1];
		
	}
	
	
	
	//打印下数组m，这个方法是多余的，只是为了输出m的值，以供我们方便用肉眼验证程序的正确性
	public static void printMatrix()
	{

		for ( int i = 0; i < M; i ++)
		{
			for ( int j = 0; j < N; j ++)
			{
				System.out.printf("%5d," , x[i]+y[j]);
			}
			System.out.println();
		}
		
	}

}

 

 

 

评论

2 楼 leonzhx 2010-10-24  

kimmking 写道

楼主的算法复杂度是 M*N

不太理解，能解释一下么？

我的计算如下：

外层循环：

for (int i = 0  ; i < MIN ; i ++ )   

总共 min(M,N)次

内层循环

while (j < MAX && index < K)   

加

while ( index < K )   
            {   
            if (j < MAX)   
                temp[index ++] = x[i] + y[j ++];   
            else  
                temp[index ++] = result[k ++];   
            }   

循环 K次

所以是 O(K*min(M,N) )

1 楼 kimmking 2010-10-23  

楼主的算法复杂度是 M*N