switchboard-utils 0.10.0

use rust_decimal::Decimal;
use std::cmp::Ordering;

pub trait MedianValue:
    Copy + std::ops::Add<Output = Self> + std::ops::Div<Output = Self> + PartialOrd
{
    fn two() -> Self;
}

impl MedianValue for Decimal {
    fn two() -> Self {
        Decimal::TWO
    }
}

impl MedianValue for i128 {
    fn two() -> Self {
        2
    }
}

impl MedianValue for u64 {
    fn two() -> Self {
        2
    }
}

pub fn median<T: MedianValue>(values: &[T]) -> Option<T> {
    if values.is_empty() {
        return None;
    }

    let mut values = values.to_vec();
    // Use `sort_unstable_by` to avoid allocating a new vector and using more memory. Because we
    // don't care about the relative order of equal elements, this is safe.
    values.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal));
    let mid = values.len() / 2;
    if values.len() % 2 == 0 {
        Some((values[mid - 1] + values[mid]) / T::two())
    } else {
        Some(values[mid])
    }
}

pub fn weighted_median<T: MedianValue>(values: &[(T, u32)]) -> Option<T> {
    let mut values = values.to_vec();
    // Use `sort_unstable_by` to avoid allocating a new vector and using more memory. Because we
    // don't care about the relative order of equal elements, this is safe.
    values.sort_unstable_by(|(a, _), (b, _)| a.partial_cmp(b).unwrap_or(Ordering::Equal));

    // Calculate total weight
    let total_weight: u32 = values.iter().map(|(_, w)| w).sum();
    // Add 1 to ensure we have strictly >50% of total weight, handling both odd/even cases
    let needed_weight = (total_weight + 1) / 2;

    // Find weighted median
    let mut cumulative_weight = 0u32;
    for (value, weight) in values.iter() {
        cumulative_weight += weight;
        if cumulative_weight >= needed_weight {
            return Some(*value);
        }
    }

    // Fallback to None if we somehow didn't return earlier
    None
}

#[cfg(test)]
mod tests {
    use super::{median, weighted_median};

    mod decimal {
        use super::*;
        use rust_decimal::Decimal;
        use rust_decimal_macros::dec;

        #[test]
        fn odd_length() {
            let values = vec![dec!(1.0), dec!(3.0), dec!(2.0)];
            assert_eq!(median(&values), Some(dec!(2.0)));
        }

        #[test]
        fn even_length() {
            let values = vec![dec!(1.0), dec!(2.0), dec!(3.0), dec!(4.0)];
            assert_eq!(median(&values), Some(dec!(2.5)));
        }

        #[test]
        fn empty() {
            let values: Vec<Decimal> = vec![];
            assert_eq!(median(&values), None);
        }

        #[test]
        fn single_element() {
            let values = vec![dec!(1.0)];
            assert_eq!(median(&values), Some(dec!(1.0)));
        }

        #[test]
        fn unsorted() {
            let values = vec![dec!(5.0), dec!(2.0), dec!(1.0), dec!(3.0), dec!(4.0)];
            assert_eq!(median(&values), Some(dec!(3.0)));
        }

        #[test]
        fn negative_numbers() {
            let values = vec![dec!(-5.0), dec!(-2.0), dec!(1.0), dec!(3.0), dec!(4.0)];
            assert_eq!(median(&values), Some(dec!(1.0)));
        }
    }

    mod weighted_decimal {
        use super::*;
        use rust_decimal::Decimal;
        use rust_decimal_macros::dec;

        #[test]
        fn basic_weighted() {
            let values = vec![(dec!(1.0), 1), (dec!(2.0), 2), (dec!(3.0), 1)];
            assert_eq!(weighted_median(&values), Some(dec!(2.0)));
        }

        #[test]
        fn heavy_weight_dominates() {
            let values = vec![(dec!(1.0), 1), (dec!(5.0), 10), (dec!(2.0), 1)];
            assert_eq!(weighted_median(&values), Some(dec!(5.0)));
        }

        #[test]
        fn empty() {
            let values: Vec<(Decimal, u32)> = vec![];
            assert_eq!(weighted_median(&values), None);
        }

        #[test]
        fn single_element() {
            let values = vec![(dec!(1.0), 5)];
            assert_eq!(weighted_median(&values), Some(dec!(1.0)));
        }

        #[test]
        fn unsorted() {
            let values = vec![(dec!(5.0), 1), (dec!(2.0), 3), (dec!(1.0), 1)];
            assert_eq!(weighted_median(&values), Some(dec!(2.0)));
        }

        #[test]
        fn negative_numbers() {
            let values = vec![
                (dec!(-5.0), 1),
                (dec!(-2.0), 2),
                (dec!(1.0), 1),
                (dec!(3.0), 1),
                (dec!(4.0), 1),
            ];
            assert_eq!(weighted_median(&values), Some(dec!(-2.0)));
        }

        #[test]
        fn mixed_negative_with_heavy_weights() {
            let values = vec![
                (dec!(-5.0), 1),
                (dec!(-2.0), 1),
                (dec!(1.0), 10), // Heavy weight on positive number
                (dec!(3.0), 1),
            ];
            assert_eq!(weighted_median(&values), Some(dec!(1.0)));
        }
    }

    mod i128 {
        use super::*;

        #[test]
        fn odd_length() {
            let values = vec![1i128, 3i128, 2i128];
            assert_eq!(median(&values), Some(2i128));
        }

        #[test]
        fn even_length() {
            let values = vec![1i128, 2i128, 3i128, 4i128];
            assert_eq!(median(&values), Some(2i128)); // Note: Integer division
        }

        #[test]
        fn empty() {
            let values: Vec<i128> = vec![];
            assert_eq!(median(&values), None);
        }

        #[test]
        fn single_element() {
            let values = vec![42i128];
            assert_eq!(median(&values), Some(42i128));
        }

        #[test]
        fn unsorted() {
            let values = vec![5i128, 2i128, 1i128, 3i128, 4i128];
            assert_eq!(median(&values), Some(3i128));
        }

        #[test]
        fn large_numbers() {
            let values = vec![i128::MAX, 0, i128::MIN];
            assert_eq!(median(&values), Some(0));
        }

        #[test]
        fn negative_numbers() {
            let values = vec![-5i128, -2i128, 2i128, 3i128, 4i128];
            assert_eq!(median(&values), Some(2i128));
        }
    }

    mod weighted_i128 {
        use super::*;

        #[test]
        fn basic_weighted() {
            let values = vec![(1i128, 1), (2i128, 2), (3i128, 1)];
            assert_eq!(weighted_median(&values), Some(2i128));
        }

        #[test]
        fn heavy_weight_dominates() {
            let values = vec![(1i128, 1), (5i128, 10), (2i128, 1)];
            assert_eq!(weighted_median(&values), Some(5i128));
        }

        #[test]
        fn empty() {
            let values: Vec<(i128, u32)> = vec![];
            assert_eq!(weighted_median(&values), None);
        }

        #[test]
        fn single_element() {
            let values = vec![(42i128, 5)];
            assert_eq!(weighted_median(&values), Some(42i128));
        }

        #[test]
        fn unsorted() {
            let values = vec![(5i128, 1), (2i128, 3), (1i128, 1)];
            assert_eq!(weighted_median(&values), Some(2i128));
        }

        #[test]
        fn negative_numbers() {
            let values = vec![(-5i128, 1), (-2i128, 2), (1i128, 1), (3i128, 1), (4i128, 1)];
            assert_eq!(weighted_median(&values), Some(-2i128));
        }

        #[test]
        fn mixed_negative_with_heavy_weights() {
            let values = vec![
                (-5i128, 1),
                (-2i128, 1),
                (1i128, 10), // Heavy weight on positive number
                (3i128, 1),
            ];
            assert_eq!(weighted_median(&values), Some(1i128));
        }

        #[test]
        fn extreme_numbers() {
            let values = vec![(i128::MIN, 1), (0i128, 5), (i128::MAX, 1)];
            assert_eq!(weighted_median(&values), Some(0i128));
        }
    }

    mod u64 {
        use super::*;

        #[test]
        fn odd_length() {
            let values = vec![1u64, 3u64, 2u64];
            assert_eq!(median(&values), Some(2u64));
        }

        #[test]
        fn even_length() {
            let values = vec![1u64, 2u64, 3u64, 4u64];
            assert_eq!(median(&values), Some(2u64)); // Note: Integer division
        }

        #[test]
        fn empty() {
            let values: Vec<u64> = vec![];
            assert_eq!(median(&values), None);
        }

        #[test]
        fn single_element() {
            let values = vec![42u64];
            assert_eq!(median(&values), Some(42u64));
        }

        #[test]
        fn unsorted() {
            let values = vec![5u64, 2u64, 1u64, 3u64, 4u64];
            assert_eq!(median(&values), Some(3u64));
        }

        #[test]
        fn large_numbers() {
            let values = vec![u64::MAX, 0, 1];
            assert_eq!(median(&values), Some(1));
        }
    }

    mod weighted_u64 {
        use super::*;

        #[test]
        fn basic_weighted() {
            let values = vec![(1u64, 1), (2u64, 2), (3u64, 1)];
            assert_eq!(weighted_median(&values), Some(2u64));
        }

        #[test]
        fn heavy_weight_dominates() {
            let values = vec![(1u64, 1), (5u64, 10), (2u64, 1)];
            assert_eq!(weighted_median(&values), Some(5u64));
        }

        #[test]
        fn empty() {
            let values: Vec<(u64, u32)> = vec![];
            assert_eq!(weighted_median(&values), None);
        }

        #[test]
        fn single_element() {
            let values = vec![(42u64, 5)];
            assert_eq!(weighted_median(&values), Some(42u64));
        }

        #[test]
        fn unsorted() {
            let values = vec![(5u64, 1), (2u64, 3), (1u64, 1)];
            assert_eq!(weighted_median(&values), Some(2u64));
        }

        #[test]
        fn extreme_numbers() {
            let values = vec![(0u64, 1), (u64::MAX / 2, 5), (u64::MAX, 1)];
            assert_eq!(weighted_median(&values), Some(u64::MAX / 2));
        }
    }
}