Conditionally Interleave/Intersperse Iterators

I was solving today's leetcode problem, here's a simple example:

Input: nums = [-4,-1,0,3,10], (the input array is always sorted)
Output: [0,1,9,16,100]
Explanation: After squaring, the array becomes [16,1,0,9,100].
After sorting, it becomes [0,1,9,16,100].

Trying to solve it in linear time, I realized I would have to combine two slices/iterators based on item magnitudes. I tried looking for an ergonomic iterator-combining function, however, having found nothing I ended up writing it myself. Here's conditional_interleave and how I used it to solve the problem:

impl Solution {
    pub fn sorted_squares(nums: Vec<i32>) -> Vec<i32> {
        let partition_point = nums.partition_point(|&x| x < 0);
        let negative = nums[0..partition_point].iter().rev();
        let positive = nums[partition_point..].iter();

        negative
            .conditional_interleave(positive, |a, b| a.abs() < b.abs())
            .map(|x| x * x)
            .collect()
    }
}

See implementation in the (Playground).

I am interested whether:

  • there is a standard way to achieve the same result,
  • and whether this is a good implementation of the idea.

Finally, (I may be reaching here), would this be a good addition to the standard library?
I often need something like this. Another use case that comes to mind is merging sub-arrays in standard merge sort.

What you call conditional_interleave is usually called merge and is the thing that powers merge sort. It is available in the itertools crate, and the implementation is quite similar, but is uses the Peekable trait.

Here's how I'd solve it with itertools:

pub fn sorted_squares(nums: Vec<i32>) -> Vec<i32> {
    let partition_point = nums.partition_point(|&x| x < 0);
    let negative = nums[0..partition_point].iter().rev().map(|x| x * x);
    let positive = nums[partition_point..].iter().map(|x| x * x);
    negative.merge(positive).collect()
}

You may also be interested in Iterator::partition_in_place: Iterator in core::iter - Rust

4 Likes

The itertools implementation you linked seems to be using a homecooked method, called PutBack, instead of the standard Peekable. A good find though, thanks!

1 Like