"Trim" two-dimensional grid

Suppose I have a 2D grid: Vec<Vec<Item>> with

enum Item {
    A,
    B,
    // more possible fields
}

in which each row, i.e. Vec<Item> has the same length. For example:

use Item::*;
let grid = vec![
    vec![A, A, A, B, B, B, A],
    vec![A, B, A, B, B, A, A],
];

I want to "trim" Item::A from the front and back of each row in order to have no full Item::A columns there. Continuing the example from above:

assert_eq!(
    trim_grid(&grid, |item| matches!(item, Item::A)),
    vec![vec![A, A, B, B, B], vec![B, A, B, B, A],]
)

Here is my implementation

fn trim_grid<T, P>(grid: Vec<Vec<T>>, predicate: P) -> Vec<Vec<T>>
where
    P: Fn(&T) -> bool,
{
    let width = grid[0].len();
    let skip = grid
        .iter()
        .map(|row| {
            row.iter()
                .position(|item| !predicate(item))
                .unwrap_or(width)
        })
        .min()
        .unwrap();
    let take = width
        - skip
        - grid
            .iter()
            .map(|row| {
                row.iter()
                    .rev()
                    .position(|item| !predicate(item))
                    .unwrap_or(width)
            })
            .min()
            .unwrap();
    grid.into_iter()
        .map(|row| row.into_iter().skip(skip).take(take).collect())
        .collect()
}

Playground

I'm basically going through each row once from the front and once from the back to find the position of the first and last item that is not Item::A. Using the minimum across all rows, I'm able to compute the parameters for .skip() and .take() to build a new grid.

Two things bother me here:

  1. The iteration to compute skip and take are basically the same, with the former using row.iter() and the latter using row.iter().rev(). Any tips how to reduce the duplication here?
  2. Instead of removing the items from the grid in place, I'm building a new one. I would prefer to take &mut grid and avoid that. I got it working, but it was basically unintelligible.

One thing I'll note here is that you have ownership of the input grid, so you probably don't want to collect new rows, but instead trim the existing ones.

My rewrite for it ended up here: https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=b9683abee5d2ceb76b90ca9ad7bb2210

fn trim_grid<T, P>(mut grid: Vec<Vec<T>>, predicate: P) -> Vec<Vec<T>>
where
    P: Fn(&T) -> bool,
{
    let (mut skip, mut rskip) = (usize::MAX, usize::MAX);
    for row in &grid {
        let this_skip = row.iter().take_while(|x| predicate(x)).count();
        skip = skip.min(this_skip);
        let this_rskip = row.iter().rev().take_while(|x| predicate(x)).count();
        rskip = rskip.min(this_rskip);
    }
    for row in &mut grid {
        row.truncate(row.len() - rskip);
        row.drain(..skip);
    }
    grid
}

I originally tried using position & rposition, but that needed annoying off-by-one fixes for how I ended up using it, as well as dealing with None, so the "how many should I remove from each end?" approach seemed to go cleaner.

(I also wrote that first loop as a fold at first, but YMMV on whether you like that better. It's not clear that it was more readable that way.)

1 Like

That looks nice, thanks! I was aware of drain, but missed truncate.

IMO, the version without fold is easier to read.