Find Kth Smallest Element in Two Sorted Arrays

Problem

Given K and two sorted arrays with length M and N, find the K-th smallest element among all elements in the two arrays.

 

Code

import sys
class Solution(object):
    def findKthSmallest(self, nums1, nums2, k):
        """
        :type nums1: sorted List[int]
        :type nusm2: sorted List[int]
        :type k: int
        :rtype: int
        """
        if not (0 <k <= len(nums1)+len(nums2)):
            return None

        while True:
            i = int(float(len(nums1))/(len(nums1)+len(nums2))*(k-1))
            j = k-1-i
            nums1_i = nums1[i] if i < len(nums1) else sys.maxint 
            nums1_i_1 =  nums1[i-1] if i-1>=0 else -sys.maxint      
            nums2_j = nums2[j] if j < len(nums2) else sys.maxint
            nums2_j_1 = nums2[j-1] if j-1>=0 else -sys.maxint

            if nums2_j_1 <= nums1_i <= nums2_j:
                return nums1_i
            if nums1_i_1 <=nums2_j<=nums1_i:
                return nums2_j
        
            if nums1_i < nums2_j:
                k = k - i - 1
                nums1 = nums1[i+1:]
                nums2 = nums2[:j]
            else:
                k = k - j - 1
                nums1 = nums1[:i]
                nums2 = nums2[j+1:]
                

s = Solution()
print s.findKthSmallest([2,3,4,5], [], 3)
print s.findKthSmallest([], [2,3,4,5], 3)
print s.findKthSmallest([2,3,4,5,5], [6,7,8], 6)

 

Idea

A very straightforward method requires O(K) time: You can use two pointers pointing from the heads of the two lists. Every time you compare which pointer points to a smaller number and push the smaller number to your result list. After that, you make the pointer point to the next element. 

Our method can achieve O(log(M+N)) time: the idea can be found here. It is similar to the idea to find median in two sorted arrays. The difference is that at line 14 we don’t simply assign `i` with `(len(nums1)+len(nums2))/2`. We must make sure i will be assigned valid value (between 0 and len(nums1)) every time. So we assign `i` as proportionally to nums1’s length: i = int(float(len(nums1))/(len(nums1)+len(nums2))*(k-1))