6
\$\begingroup\$

I was solving this task https://leetcode.com/explore/learn/card/fun-with-arrays/525/inserting-items-into-an-array/3253/

The main idea is to merge two sorted arrays (more details in the site) Example:

Input: nums1 = [1,2,3,0,0,0], m = 3, nums2 = [2,5,6], n = 3
Output: [1,2,2,3,5,6]
Explanation: The arrays we are merging are [1,2,3] and [2,5,6].
The result of the merge is [1,2,2,3,5,6] with the underlined elements coming from nums1.

My solution is working and fine in terms of time complexity, but I thin that it can be improved:

  1. Can I avoid i1 as usize in index calculation?
  2. What rust features I can apply?

My solution (passed all tests)

impl Solution {
    pub fn merge(nums1: &mut Vec<i32>, m: i32, nums2: &mut Vec<i32>, n: i32) {
        let mut insert_pos: i32 = m + n - 1;
        let mut i1 = m - 1;
        let mut i2 = n - 1;

        while i1 >= 0 && i2 >= 0 {
            if nums1[i1 as usize] > nums2[i2 as usize] {
                nums1[insert_pos as usize] = nums1[i1 as usize];
                i1 -= 1;
            } else {
                nums1[insert_pos as usize] = nums2[i2 as usize];
                i2 -= 1;
            }
            insert_pos -= 1;
        }
        while i1 >= 0 {
            nums1[insert_pos as usize] = nums1[i1 as usize];
            i1 -= 1;
            insert_pos -= 1;
        }
        while i2 >= 0 {
            nums1[insert_pos as usize] = nums2[i2 as usize];
            i2 -= 1;
            insert_pos -= 1;
        }
    }
}
\$\endgroup\$
1
  • 1
    \$\begingroup\$ Biggest issue is that the while i1 >= 0 loop does nothing and can just be removed! See my answer for additional clean up and style suggestions. \$\endgroup\$ Commented Apr 22, 2023 at 4:14

3 Answers 3

7
\$\begingroup\$

A lot of room for improvement here! :)

For the i as usize issue, the general principle is: whenever if you have variables indexing into an array, make them usize. So you have two Vec of i32, but the indexes (and lengths) should be usize. That's m, n, insert_pos, i1, and i2.

Once you make this change, you can get rid of all those pesky as usize annotations, but now you have a different problem: you are using -1 to indicate when there's no more elements to insert, and -1 isn't a valid usize. Instead, best practice here is to add 1 to all your variables. The >=0 checks become >0 checks instead, and the -1 calculations are moved into the while loop. With this change, the beginning of the function becomes:

pub fn merge(nums1: &mut Vec<i32>, m: usize, nums2: &mut Vec<i32>, n: usize) {
    let mut insert_pos = m + n;
    let mut i1 = m;
    let mut i2 = n;

    while i1 > 0 && i2 > 0 {
        if nums1[i1 - 1] > nums2[i2 - 1] {
        ...

Our next observation for improvement is that insert_pos is always just equal to i1 + i2. So let's just replace all instances of it with i1 + i2. Once we do this, you'll notice something fascinating -- your second while loop actually does nothing and can be removed:

    while i1 > 0 {
        i1 -= 1;
        nums1[i1] = nums1[i1];
    }

Once we remove that, we're left with just the first and third loops. Before we get back to the final updated code, let's take a look at a couple of other improvements:

  • Always try running cargo clippy! It's an automatic tool which always gives helpful suggestions to make your code more Rust-like. Here, it wants us to convert the &mut Vec<i32> arguments to &mut [i32], because the vectors are not being resized at all in the function body, which is correct.

  • In fact, we can go further and remove the &mut keyword for nums2, as nums2 is only read from and not written to.

  • Finally, the variables m and n are immediately assigned to i1 and i2, and otherwise removed. So it would be cleaner to just declare i1 and i2 as mut in the function signature. That is, we can write pub fn merge(nums1: &mut [i32], mut i1: usize, ... and omit the variables m and n. I would personally also rename these to n1 and n2 for clarity.


Here's our final code, with all of the above changes:

pub fn merge(nums1: &mut [i32], mut n1: usize, nums2: &[i32], mut n2: usize) {
    while n1 > 0 && n2 > 0 {
        if nums1[n1 - 1] > nums2[n2 - 1] {
            n1 -= 1;
            nums1[n1 + n2] = nums1[n1];
        } else {
            n2 -= 1;
            nums1[n1 + n2] = nums2[n2];
        }
    }
    while n2 > 0 {
        n2 -= 1;
        nums1[n2] = nums2[n2];
    }
}

There are definitely ways to do this more concisely using more advanced Rust built-ins -- like iterators and ranges -- but not a bad start for purely imperative code.

\$\endgroup\$
0
6
\$\begingroup\$

Caleb Stanford already covered how you could improve the type of loop you wrote, but since you asked about other Rust features and how you might use them here, I’ll cover those.

One alternative would be to convert both arrays (or for more flexibility, array slices) into iterators with code like let mut m_it = m.iter();. Then, you might test the iterator with a pattern match, using if let statements, until both yield None. Then, both arrays are fully copied. There are other approaches, such as the iter::gt() method to test whether one iterator is currently greater than the other.

You could do this in a while loop, or a loop terminated by a return statement, or a tail-recursive function (although the latter is not very idiomatic in Rust, which lacks guaranteed tail-call optimization). Instead of iterators, you could also iterate over array slices, by removing the first element.

Note that most real-world merge sort algorithms actually want a merge function that can merge two array slices into a third slice in memory. This allows the algorithm to re-use a pair of buffers, ping-ponging between them at every step of the algorithm, without doing any additional memory management.

So, this (intentionally) doesn’t solve the LeetCode problem, which you want to do on your own, but it does serve as a demonstration of how to write a merge function with iterators. It’s not especially optimized. (This implementation doesn’t use the known sizes of the slices to help the optimizer out.)

pub fn merge<'out>( xs: &[i32],
                    ys: &[i32],
                    outputs: &'out mut[i32]
            ) -> &'out mut [i32] {
    assert!(xs.len() + ys.len() == outputs.len());

    let mut x_it = xs.iter().peekable();
    let mut y_it = ys.iter().peekable();
    let mut out_it = outputs.iter_mut();

    loop { // The terminating case is at the bottom:
        if let Some(x) = x_it.peek() {
            if let Some(y) = y_it.peek() {
                /* Unwrapping is safe because we already looked ahead at x_it
                 * and y_it, and tested that outputs is the same size as the
                 * inputs combined.
                 */

                 *out_it.next().unwrap() = if y > x {
                    *x_it.next().unwrap()
                 } else {
                    *y_it.next().unwrap()
                 };
            } else { // Remaining xs but no ys.
                *out_it.next().unwrap() = *x_it.next().unwrap();
            }
        } else if let Some(&y) = y_it.next() { // Remaining ys but no xs
            *out_it.next().unwrap() = y;
        } else {
             return outputs; // Terminating case.
        } // end if let
    } // end loop
} // end fn

Rust’s type system makes us get really pedantic, but it can deduce a lot of things itself, and there is a reason for all of it. In this example, we need to create Peekable input iterators, to be able to test and compare their values without advancing the iterators. Output iterators always need to borrow the data they reference mutably. The iterators themselves must also be mut, so we can advance them. Calling next or peek on them returns an Option, and since the iterator is to a borrowed container of i32, it’s an Option<&i32> and we need to dereference the value inside Some. Finally, I chose (unlike LeetCode) to return the output slice, which is very useful to chain and compose merge operations. But, any reference we return in Rust needs a lifetime. In this case, we’re returning the output slice that we received, so we give its lifetime a name and use that. Most people give lifetimes names like <'a, 'b>, but if I wouldn’t name a variable that way, I don’t name a lifetime that way either.

You might flatten the pyramid of doom here a little with a match statement, such as:

pub fn merge<'out>( xs: &[i32],
                    ys: &[i32],
                    outputs: &'out mut[i32]
            ) -> &'out mut [i32] {
    let mut x_it = xs.iter().peekable();
    let mut y_it = ys.iter().peekable();
    let mut out_it = outputs.iter_mut();

    loop { // The terminating case is at the top:
        match (x_it.peek(), y_it.peek()) {
            (None, None) => {
                assert!(out_it.next().is_none());
                return outputs;
            },
            (Some(_), None) => {*out_it.next().unwrap() = *x_it.next().unwrap()},
            (None, Some(_)) => {*out_it.next().unwrap() = *y_it.next().unwrap()},
            (Some(x), Some(y)) => {
                *out_it.next().unwrap() = if x > y {
                    *y_it.next().unwrap()
                } else {
                    *x_it.next().unwrap()
                };
            }
        }
    } // end loop
} // end fn

One way to improve this a bit is not to repeat the checks that either input is empty once you know it is. At that point, you can just copy the remaining items using a while let loop.

pub fn merge<'out>( xs: &[i32],
                    ys: &[i32],
                    outputs: &'out mut[i32]
            ) -> &'out mut [i32] {
    assert!(xs.len() + ys.len() == outputs.len());

    let mut x_it = xs.iter().peekable();
    let mut y_it = ys.iter().peekable();
    let mut out_it = outputs.iter_mut();

    loop { // The terminating case is at the top:
        match (x_it.peek(), y_it.peek()) {
            (None, None) => {
                assert!(out_it.next().is_none());
                return outputs;
            },
            (Some(_), None) => {
                while let Some(&x) = x_it.next() {
                    *out_it.next().unwrap() = x;
                }
            },
            (None, Some(_)) => {
                while let Some(&y) = y_it.next() {
                    *out_it.next().unwrap() = y;
                }
            },
            (Some(x), Some(y)) => {
                *out_it.next().unwrap() = if x > y {
                    *y_it.next().unwrap()
                } else {
                    *x_it.next().unwrap()
                };
            },
        }
    } // end loop
} // end fn

You could instead do this with ranges and a conditional block:

pub fn merge<'out>( xs: &[i32],
                    ys: &[i32],
                    outputs: &'out mut[i32]
            ) -> &'out mut [i32] {
    assert!(xs.len() + ys.len() == outputs.len());

    let mut x_idx : usize = 0;
    let mut y_idx : usize = 0;

/* Since this loop maintains the invariant that x_idx + y_idx == i, and we
 * already tested that xs_len() + ys_len() == outputs.len(), x_idx and y_idx
 * cannot both point to the end of the slice while i < outputs.len().
 */
    for i in 0..outputs.len() {
        outputs[i] =
            if y_idx == ys.len() {
                let x = xs[x_idx];
                x_idx += 1;
                x
            } else if x_idx == xs.len() {
                let y = ys[y_idx];
                y_idx += 1;
                y
            } else if xs[x_idx] > ys[y_idx] {
                let y = ys[y_idx];
                y_idx += 1;
                y 
            } else {
                let x = xs[x_idx];
                x_idx += 1;
                x
            };
    }

    outputs
} // end fn

This ends up checking the bounds of the slices twice, once when testing the conditions on the indices, and again when dereferencing with them. You could therefore optimize this a bit more if you are willing to use unsafe code.

The optimizer generally does well with code like this, where it can see that it is iterating over an array slice in order and setting each element to the same conditional expression (such as a lookup from one of two other slices). In this case, the conditional is a bit too complicated for rustc to be able to generate vectorized code, but many other loops similar to this one auto-vectorize. (The unsafe version optimizes well enough that rustc can optimize out the call to merge on two constant arrays, and pre-calculate the result at compile time.)

Or, you can combine the idea of iterating over the output elements in order with input iterators. The code this generates also optimizes well, and is more explicit about covering all possible cases.

pub fn merge<'out>( xs: &[i32],
                    ys: &[i32],
                    outputs: &'out mut[i32]
            ) -> &'out mut [i32] {
    assert!(xs.len() + ys.len() == outputs.len());

    let mut x_it = xs.into_iter().peekable();
    let mut y_it = ys.into_iter().peekable();

    outputs.iter_mut()
           .for_each(move|p| {
        *p = match (x_it.peek(), y_it.peek()) {
            (Some(&&x), None) => {
                x_it.next();
                x
            },
            (None, Some(&&y)) => {
                y_it.next();
                y
            },
            (Some(&&x), Some(&&y)) => {
                if x > y {
                     y_it.next();
                     y
                } else {
                     x_it.next();
                     x
                }
            },
            (None, None) => {
 /* Since we have already tested that there are the same number of output
  * as input elements, it is impossible to consume all the inputs before all of
  * the outputs.
  */
                unreachable!()
            },
        };
            });

    outputs
} // end fn

Unlike the version using a for loop and the if/else if blocks, where I just left out the check for both iterators being empty and wrote a comment explaining why that should never happen, the match statement forces me to cover every possible pattern. The unreachable!() macro is very handy for logic errors such as this, and will make the code fail fast if it somehow occurs.

And a simple test driver:

pub fn main() {
    assert_eq!(merge(&[], &[], &mut[]), []);
    assert_eq!(merge(&[1,2,3], &[], &mut[0;3]), [1,2,3]);
    assert_eq!(merge(&[], &[4,5,6], &mut[0;3]), [4,5,6]);
    assert_eq!(merge(&[1,3,5], &[2,4,6], &mut[0;6]), [1,2,3,4,5,6]);
}
\$\endgroup\$
3
\$\begingroup\$

For those looking for real-world code in addition to learning exercises, I would suggest using itertools::Itertools::merge, which implements exactly this problem based on iterators.

\$\endgroup\$
2
  • 1
    \$\begingroup\$ This feels more like a comment than an answer. Are there any insights you find particularly relevant? \$\endgroup\$
    – Davislor
    Commented Apr 24, 2023 at 5:17
  • \$\begingroup\$ This should be the preferred answer. \$\endgroup\$
    – Steven He
    Commented Jan 9 at 12:53

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.