使用函数式编程方法迭代网格的 (x,y) 位置

Question

我想迭代(x, y)二维网格中的位置并为每个位置调用一个函数（提供x和y作为参数）。我知道如何使用循环来完成此操作for，但我想使用函数式编程方法来编写此操作，以便稍后可以利用 Rayon 等库。我已经成功构建了一个功能版本，但它看起来很复杂，我想问是否有更干净的方法来实现这一点。

下面是计算网格中所有 2x2 面片的最大总和的小示例：

use std::cmp::max;

const WIDTH: usize = 4;
const HEIGHT: usize = 3;

type Grid = [[u32; WIDTH]; HEIGHT];

fn main() {
    let grid: Grid = [
        [1, 3, 5, 8], 
        [3, 9, 4, 2], 
        [3, 4, 5, 0],
    ];
                          
    let coords = (0..WIDTH - 1).flat_map(|x| (0..HEIGHT - 1).map(move |y| (x, y)));
    let max_sum = coords.map(|(x, y)| sum_2x2(x, y, &grid)).max().unwrap();
    println!("Max 2x2 patch: {}", max_sum);
}

fn sum_2x2(x: usize, y: usize, grid: &Grid) -> u32 {
    [
        grid[y][x],
        grid[y][x + 1],
        grid[y + 1][x],
        grid[y + 1][x + 1],
    ]
    .iter()
    .sum()
}

对于这样一个简单的任务来说，这条线let coords = (0..WIDTH - 1).flat_map(|x| (0..HEIGHT - 1).map(move |y| (x, y)));似乎相当复杂。在Python中，我会执行以下操作来获取位置（我认为这更清晰）：

use std::cmp::max;

const WIDTH: usize = 4;
const HEIGHT: usize = 3;

type Grid = [[u32; WIDTH]; HEIGHT];

fn main() {
    let grid: Grid = [
        [1, 3, 5, 8], 
        [3, 9, 4, 2], 
        [3, 4, 5, 0],
    ];
                          
    let coords = (0..WIDTH - 1).flat_map(|x| (0..HEIGHT - 1).map(move |y| (x, y)));
    let max_sum = coords.map(|(x, y)| sum_2x2(x, y, &grid)).max().unwrap();
    println!("Max 2x2 patch: {}", max_sum);
}

fn sum_2x2(x: usize, y: usize, grid: &Grid) -> u32 {
    [
        grid[y][x],
        grid[y][x + 1],
        grid[y + 1][x],
        grid[y + 1][x + 1],
    ]
    .iter()
    .sum()
}

有没有更好的方法来写这个或者我只需要习惯它？

Answer 1

正如@hellow 在评论中指出的那样，板条箱中有一个iproduct宏itertools可以完成我想要的操作。使用它，我的代码可以重写如下：

use itertools::iproduct;
...
let coords = iproduct!(0..WIDTH - 1, 0..HEIGHT - 1);
...

这正是我一直在寻找的。感谢所有在这里评论/发布答案的人。

PS：如果您想将 itertools 迭代器与 rayon 一起使用，您可以使用par_bridge方法来完成此操作，例如iproduct!(0..WIDTH - 1, 0..HEIGHT - 1).par_iter();。我花了一段时间才弄清楚这一点，但par_iter不起作用into_par_iter。