Partition a slice in place.

partition.rs

/// Partitions a slice in place by a predicate.
///
/// The elements in `slice` for which `predicate` returns `true` are all moved to the left side of the slice, leaving
/// the elements for which `predicate` returns `false` on the right side.
///
/// If the slice can be partitioned into two non-empty parts, returns `Some` with the index of the first element of the
/// right partition. Otherwise, returns `None`.
///
/// # Examples
/// ```
/// let mut vec = vec![1, 2, 3, 4, 5];
///
/// // Partition between even and odd numbers
/// let p = partition(&mut vec, |x| x % 2 == 0);
///
/// // The exact order of elements in the result is not specified, but all even elements
/// // are before the split point and all odd elements are after the split point
/// assert_eq!(vec, [2, 4, 3, 1, 5]);
///
/// // Return value is a Some with the index of the first odd element
/// assert_eq!(p, Some(2));
///
/// let index = p.unwrap();
/// assert_eq!(vec[0..index], [2, 4]);
/// assert_eq!(vec[index..], [3, 1, 5]);
/// ```
///
/// ```
/// let mut vec = vec![64, 12, 20];
///
/// let p = partition(&mut vec, |x| x % 2 == 0);
///
/// // No odd numbers
/// assert_eq!(p, None);
/// ```
///
/// ```
/// let mut vec = vec![65, 13, 21];
///
/// let p = partition(&mut vec, |x| x % 2 == 0);
///
/// // No even numbers
/// assert_eq!(p, None);
/// ```
fn partition<E, P: Fn(&E) -> bool>(slice: &mut [E], predicate: P) -> Option<usize> {
    if slice.is_empty() {
        return None;
    }

    // https://en.wikipedia.org/wiki/Quicksort
    // Lomuto partition scheme

    let mut j = 0;
    let len = slice.len();
    for i in 0..len {
        if predicate(&slice[i]) {
            slice.swap(i, j);
            j += 1;
        }
    }

    if j > 0 && j < len {
        Some(j)
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use crate::partition;

    #[test]
    fn test_partition() {
        let mut vec = vec![1, 2, 3, 4, 5];

        // Partition between even and odd numbers
        let p = partition(&mut vec, |x| x % 2 == 0);

        // The exact order of elements in the result is not specified, but all even elements
        // are before the split point and all odd elements are after the split point
        assert_eq!(vec, [2, 4, 3, 1, 5]);

        // Return value is a Some with the index of the first odd element
        assert_eq!(p, Some(2));

        let index = p.unwrap();
        assert_eq!(vec[0..index], [2, 4]);
        assert_eq!(vec[index..], [3, 1, 5]);
    }

    #[test]
    fn test_partition_evens() {
        let mut vec = vec![64, 12, 20];

        let p = partition(&mut vec, |x| x % 2 == 0);

        // No odd numbers
        assert_eq!(p, None);
    }

    #[test]
    fn test_partition_odds() {
        let mut vec = vec![65, 13, 21];

        let p = partition(&mut vec, |x| x % 2 == 0);

        // No even numbers
        assert_eq!(p, None);
    }

    #[test]
    fn test_partition_empty() {
        let mut vec = vec![];

        let p = partition(&mut vec, |x| x % 2 == 0);

        assert_eq!(p, None);
    }
}