Kth Largest in N Arrays

Posted 2020-08-04 Sheryl Wang

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的代码如下，复杂度Ｏ(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]