Sharing my java solution with finding kth element as helper function


  • 0
    C
    public class Solution {
        public double findMedianSortedArrays(int[] nums1, int[] nums2) {
            int total = nums1.length + nums2.length;
            if ( total % 2 == 1 ){ //median in the middle number of two arrays
                return (double) findKth( nums1, 0, nums1.length-1, nums2, 0, nums2.length-1,total/2+1 );
            }else{ // median is the average of two middle numbers
                int part1 =  findKth( nums1,0,nums1.length-1,nums2,0,nums2.length-1, total/2);
                int part2 =  findKth( nums1,0,nums1.length-1,nums2,0,nums2.length-1,total/2 + 1);
                return (double) (part1+part2)/2;
            }
        }
    
        /****
         * Helper function to find the kth element of sub-ranges of two arrays
         */
        private int findKth ( int[] nums1, int start1, int end1, int[] nums2, int start2, int end2, int k ){
            if (start1 > end1) { //no elements in first array
                return nums2[start2 + k - 1];
            } else if (start2 > end2) { // no elements in second array
                return nums1[start1 + k - 1];
            } else if (k == 1) { // result must be at the head of arrays
                return nums1[start1] < nums2[start2] ? nums1[start1] : nums2[start2];
            } else if (((end1 - start1 + 1) + (end2 - start2 + 1)) == k) { // result in the tails of arrays
                return nums1[end1] > nums2[end2] ? nums1[end1] : nums2[end2];
            }
            
            // Idetnify the parts of array can be possibly cuts
            int cut1 = (end1 - start1 + 1) < k / 2 ? (end1 - start1 + 1) : k / 2;
            int cut2 = (end2 - start2 + 1) < k / 2 ? (end2 - start2 + 1) : k / 2;
            
            if (nums1[start1 + cut1 - 1] == nums2[start2 + cut2 - 1] && k == (cut1 + cut2)) {
                //result is the cutting element
                return nums1[start1 + cut1 - 1];
            } else if (nums1[start1 + cut1 - 1] == nums2[start2 + cut2 - 1] && (k + 1) == cut1 + cut2) {
                //reust is the next of cutting element
                int next1 = (start1 + cut1) < nums1.length ? nums1[start1 + cut1] : Integer.MAX_VALUE;
                int next2 = (start2 + cut2) < nums2.length ? nums2[start2 + cut2] : Integer.MAX_VALUE;
                return next1 <= next2 ? next1 : next2;
            } else if (nums1[start1 + cut1 - 1] <= nums2[start2 + cut2 - 1]) {
                //cut first part of the first array
                return findKth(nums1, start1 + cut1, end1, nums2, start2, (cut1 + cut2) >= k ? (start2 + cut2 - 1) : end2, k - cut1);
            } else { //if ( nums1[start1+cut1-1] > nums2[start2+cut2-1]){
                //cut first part of the second array
                return findKth(nums1, start1, (cut1 + cut2) >= k ? (start1 + cut1 - 1) : end1, nums2, start2 + cut2, end2, k - cut2);
            }
        }
    }
    

    P.S.
    I got asked a similar question in google interview, which is a special case of this question:
    Find median of two sorted arrays, where two arrays are equal length.
    Maybe there is simpler solution for the special case ?


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.