hard - heap

mail@pastecode.io avatar
unknown
java
a year ago
1.4 kB
0
Indexable
Never
class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int n1 = nums1.length, n2 = nums2.length;
        int i = 0;
        PriorityQueue<Integer> min = new PriorityQueue<>();
        PriorityQueue<Integer> max = new PriorityQueue<>((a, b) -> b - a);

        while (i < n1 && i < n2) {
            if (nums1[i] <= nums2[i]) {
                max.add(nums1[i]);
                min.add(nums2[i]);
            } else {
                min.add(nums1[i]);
                max.add(nums2[i]);
            }

            swap(min, max);

            ++i;
        }

        while (i < n1) {
            max.add(nums1[i]);
            if (max.size() > (min.size() + 1)) {
                min.add(max.poll());
            }

            swap(min, max);
            ++i;
        }
        while (i < n2) {
            max.add(nums2[i]);
            if (max.size() > (min.size() + 1)) {
                min.add(max.poll());
            }

            swap(min, max);
            ++i;
        }

        return max.size() > min.size() ? (double)max.peek() : ((double) (max.peek() + min.peek()) / 2);
    }

    private static void swap(PriorityQueue<Integer> min, PriorityQueue<Integer> max) {
        if (min.peek() != null && max.peek() > min.peek()) {
            var temp = min.poll();
            min.add(max.poll());
            max.add(temp);
        }
    }
}