leetcode

0004-median-of-two-sorted-arrays.rs (3593B)

/// Given two sorted arrays nums1 and nums2 of size m and n respectively,
/// return the median of the two sorted arrays.
///
/// The overall run time complexity should be O(log (m+n)).
///
///
///
/// Example 1:
///
/// Input: nums1 = [1,3], nums2 = [2]
/// Output: 2.00000
/// Explanation: merged array = [1,2,3] and median is 2.
///
/// Example 2:
///
/// Input: nums1 = [1,2], nums2 = [3,4]
/// Output: 2.50000
/// Explanation: merged array = [1,2,3,4] and median is (2 + 3) / 2 = 2.5.
///
///
///
/// Constraints:
///
///     nums1.length == m
///     nums2.length == n
///     0 <= m <= 1000
///     0 <= n <= 1000
///     1 <= m + n <= 2000
///     -106 <= nums1[i], nums2[i] <= 106
///
use leetcode::arg_into;

struct Solution {}

impl Solution {
    // Arrays a and b can be partitioned into elements which are less
    // than or equal to the merged array's median (*_l) and those which
    // are greater or equal (*_g). A constraint that that can be derived
    // from this is that all elements in the *_l partitions must be smaller
    // than or equal those in the *_g partitions:
    //
    // - a[a_l-1] <= b[b_l]
    // - b[b_l-1] <= a[a_l]
    //
    // Additionally, since the median by definition is the (average) value
    // of an array's center element(s), the partition lengths must
    // satisfy the constraint:
    //
    // - len(b_l) = (len(a) + len(b) + 1) / 2 - len(a_l)
    //
    // Using these contraints, we can binary search for the partition
    // length a_l to determine all partitions and ultimately the median
    // of the merged array.

    pub fn find_median_sorted_arrays(nums1: Vec<i32>, nums2: Vec<i32>) -> f64 {
        if nums1.is_empty() && nums2.is_empty() {
            return 0f64;
        }

        let (a, b) = if nums1.len() > nums2.len() {
            (nums1, nums2)
        } else {
            (nums2, nums1)
        };

        let mut min = 0;
        let mut max = a.len();
        let n_l = (a.len() + b.len() + 1) / 2;
        while min != max {
            let a_l = (max + min) / 2;
            if a_l > n_l || a_l > 0 && n_l - a_l < b.len() && a[a_l - 1] > b[n_l - a_l] {
                max = a_l - 1;
            } else if n_l - a_l > b.len()
                || n_l - a_l > 0 && a_l < a.len() && b[n_l - a_l - 1] > a[a_l]
            {
                min = a_l + 1;
            } else {
                min = a_l;
                max = a_l;
            }
        }

        let a_l = min as i32;
        let b_l = n_l as i32 - a_l;
        let v = |x: &Vec<i32>, i| {
            if i >= 0 {
                x.get(i as usize).cloned()
            } else {
                None
            }
        };
        let m1 = [v(&a, a_l - 1), v(&b, b_l - 1)]
            .iter()
            .filter_map(|&p| p)
            .max()
            .unwrap();
        if (a.len() + b.len()) % 2 == 1 {
            m1 as f64
        } else {
            let m2 = [v(&a, a_l), v(&b, b_l)]
                .iter()
                .filter_map(|&p| p)
                .min()
                .unwrap();
            (m1 as f64 + m2 as f64) / 2.0
        }
    }
}

pub fn main() {
    println!(
        "{:?}",
        Solution::find_median_sorted_arrays(arg_into(1), arg_into(2))
    );
}

#[cfg(test)]
mod tests {
    use crate::Solution;
    use leetcode::vi;

    #[test]
    fn test() {
        assert_eq!(
            Solution::find_median_sorted_arrays(vi("[1,3]"), vi("[2]")),
            2.0
        );
        assert_eq!(
            Solution::find_median_sorted_arrays(vi("[1,2]"), vi("[3,4]")),
            2.5
        );
    }
}