Solution to "Data Stream as Disjoint Intervals" Leetcode problem

lib.rs

use std::{
    cmp::{self, Ordering},
    collections::BTreeSet,
};

#[derive(Default)]
pub struct SummaryRanges {
    ranges: BTreeSet<ClosedRange>,
}

// A closed range.
#[derive(Debug, Clone, Copy)]
struct ClosedRange {
    begin: i32,
    end: i32,
}

impl Ord for ClosedRange {
    fn cmp(&self, other: &Self) -> Ordering {
        if self.end < other.begin.saturating_sub(1) {
            // . . a . . . b . . . . . .
            // . . . . . . . . c . . . d
            Ordering::Less
        } else if self.end == other.begin.saturating_sub(1) {
            // . . a . . . b . . . . .
            // . . . . . . . c . . . d
            Ordering::Equal
        } else if self.begin > other.end.saturating_add(1) {
            // . . . . . . . . a . . . b
            // . . c . . . d . . . . . .
            Ordering::Greater
        } else if self.begin == other.end.saturating_add(1) {
            // . . . . . . . a . . . b
            // . . c . . . d . . . . .
            Ordering::Equal
        } else {
            // Overlapping means equality.
            Ordering::Equal
        }
    }
}

impl PartialOrd for ClosedRange {
    fn partial_cmp(&self, other: &ClosedRange) -> Option<cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Eq for ClosedRange {}

impl PartialEq for ClosedRange {
    fn eq(&self, other: &ClosedRange) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}

impl ClosedRange {
    /// Creates a "range" that contains only one value.
    pub fn new(value: i32) -> Self {
        ClosedRange {
            begin: value,
            end: value,
        }
    }

    /// Merges two ranges together. The ranges must overlap or share a border.
    pub fn merge(self, other: ClosedRange) -> ClosedRange {
        debug_assert_eq!(self, other);
        let begin = cmp::min(self.begin, other.begin);
        let end = cmp::max(self.end, other.end);
        ClosedRange { begin, end }
    }

    /// Converts the range into a pair of numbers.
    pub fn to_vec(&self) -> Vec<i32> {
        vec![self.begin, self.end]
    }
}

impl SummaryRanges {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn add_num(&mut self, val: i32) {
        let mut temp_range = ClosedRange::new(val);
        while let Some(range) = self.ranges.take(&temp_range) {
            temp_range = range.merge(temp_range);
        }
        self.ranges.insert(temp_range);
    }

    pub fn get_intervals(&self) -> Vec<Vec<i32>> {
        self.ranges.iter().map(ClosedRange::to_vec).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn range() {
        let one = ClosedRange::new(1);
        let two = ClosedRange::new(2);
        let three = ClosedRange::new(3);

        assert_eq!(one.cmp(&two), Ordering::Equal);
        assert_eq!(two.cmp(&one), Ordering::Equal);

        assert_ne!(one.cmp(&three), Ordering::Equal);
        assert_ne!(three.cmp(&one), Ordering::Equal);

        let one_two = one.merge(two);
        assert_eq!(one_two.begin, 1);
        assert_eq!(one_two.end, 2);

        assert_eq!(one_two.cmp(&three), Ordering::Equal);
        assert_eq!(three.cmp(&one_two), Ordering::Equal);

        let one_three = one_two.merge(three);
        assert_eq!(one_three.begin, 1);
        assert_eq!(one_three.end, 3);

        assert_eq!(one_three.cmp(&two), Ordering::Equal);
        assert_eq!(two.cmp(&one_three), Ordering::Equal);
    }

    #[test]
    fn it_works() {
        let mut summary = SummaryRanges::new();

        assert!(summary.get_intervals().is_empty());

        summary.add_num(1);
        assert_eq!(summary.get_intervals(), vec!(vec!(1, 1)));

        summary.add_num(3);
        assert_eq!(summary.get_intervals(), vec!(vec!(1, 1), vec!(3, 3)));

        summary.add_num(7);
        assert_eq!(
            summary.get_intervals(),
            vec!(vec!(1, 1), vec!(3, 3), vec!(7, 7))
        );

        summary.add_num(2);
        assert_eq!(summary.get_intervals(), vec!(vec!(1, 3), vec!(7, 7)));

        summary.add_num(6);
        assert_eq!(summary.get_intervals(), vec!(vec!(1, 3), vec!(6, 7)));

        summary.add_num(6);
        assert_eq!(summary.get_intervals(), vec!(vec!(1, 3), vec!(6, 7)));


        summary.add_num(2);
        assert_eq!(summary.get_intervals(), vec!(vec!(1, 3), vec!(6, 7)));
    }
}