Kth Largest in N Arrays

Find K-th largest element in N arrays.

Example

In n=2 arrays [[9,3,2,4,7],[1,2,3,4,8]], the 3rd largest element is 7.

In n=2 arrays [[9,3,2,4,8],[1,2,3,4,2]], the 1st largest element is 9, 2nd largest element is 8, 3rd largest element is 7 and etc.

这题很有意思,一开始拿到手,完全没有头绪.但是后来想了下其实这题是find kth largest element的follow up,只是把一个数组拓展为一个数组群,每次在数组群中查找.所以加上一个具体判断数组的处理之后,使用

partition或者heap的方法都可以.而判断当前处理到第几个数组的处理办法是维护一个数组长度的prefix sum.如果当前处理到下标已经大于当前数组的prefixsum,则移动到下一个数组.需要注意的是空数组的规避.如果数组为空,则需要往后查找.

用heap的代码如下,复杂度O(nlogk),n为总元素的数目:

class Solution:
    # @param {int[][]} arrays a list of array
    # @param {int} k an integer
    # @return {int} an integer, K-th largest element in N arrays
    def KthInArrays(self, arrays, k):
        # follow up of find K largest numbers in array, partition or heap solution 
        #the only difference is to calculate the index of number in the really array
        import heapq
        if not arrays:
            return -1 
        n = len(arrays)
        preLen = [0] #数组长度的prefix sum处理
        for i in xrange(n):
            preLen.append(preLen[-1]+len(arrays[i]))
        arrayIndex = 0
        heap = []
        i = 0
        while arrayIndex < len(preLen) - 1:
            if arrays[arrayIndex]: 非常关键,规避空数组
                if len(heap) < k:
                    heapq.heappush(heap, arrays[arrayIndex][i - preLen[arrayIndex]])
                else:
                    heapq.heappushpop(heap, arrays[arrayIndex][i - preLen[arrayIndex]])
                i += 1
            if  i == preLen[arrayIndex+1]:
                arrayIndex += 1
        
        return heap[0]

 

posted on 2016-08-27 16:33  Sheryl Wang  阅读(296)  评论(0编辑  收藏  举报

导航