The cleanest findKth (kth of two sorted array) -- try to prove me wrong, thx!


  • 0
    Q
    template <typename RandomIt, typename T = typename std::iterator_traits<RandomIt>::value_type>
    T findKth(RandomIt first1, RandomIt last1, RandomIt first2, RandomIt last2, size_t k) {
        // k is 0-based, the key point is choose two prefix sub-array, and let the sum of
        // length as large as possile, -- up to k + 1 -- and skip the smaller one.
        size_t n1 = last1 - first1, n2 = last2 - first2;
        ++k;
        for (size_t a1, a2; k > 1 && n1 > 0 && n2 > 0; ) {
            a1 = std::min(k / 2, n1);
            a2 = std::min(k - a1, n2);
            if (first1[a1-1] <= first2[a2-1]) { first1 += a1, n1 -= a1, k -= a1; }
            else { first2 += a2, n2 -= a2, k -= a2; }
        }
        --k;
        if (n1 == 0) return first2[k];
        if (n2 == 0) return first1[k];
        return std::min(*first1, *first2);
    }
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int sz1 = nums1.size(), sz2 = nums2.size();
        int half = (sz2 + sz1) / 2;
        double res = findKth(nums1.begin(), nums1.end(), nums2.begin(), nums2.end(), half);;
        if ((sz1 & 1) == (sz2 & 1)) res = (res + findKth(nums1.begin(), nums1.end(), nums2.begin(), nums2.end(), half-1)) / 2;
        return res;
    }

Log in to reply
 

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