如何用Rust复刻JS中支持任意深度的多数组求和函数

先看JavaScript中实现该功能的示例函数：

function add(...args){
  if (args.length === 0) return [];
  const a = args[0];
  let res = [];
  for (let i = 0; i < a.length; i++){
    const elements = args.map(x => x[i]);
    res.push(Array.isArray(a[i]) ? add(...elements) : elements.reduce((a, b) => a + b));
  }
  return res;
}

console.log(add([[1]], [[2]])); // [[3]]
console.log(add([1], [2])); // [3]
console.log(add([1, 2], [3, 4])); // [4, 6]
console.log(add([[1, 2]], [[3, 4]])); // [[4, 6]]

Rust实现方案 #

由于Rust是静态类型语言，我们先定义枚举统一表示数字或嵌套数组结构，再实现递归求和逻辑：

#[derive(Debug, Clone, PartialEq)]
enum NestedNum {
    Number(i32),
    Array(Vec<NestedNum>),
}

fn add(arrays: &[Vec<NestedNum>]) -> Vec<NestedNum> {
    if arrays.is_empty() {
        return Vec::new();
    }

    let first_array = &arrays[0];
    let mut result = Vec::with_capacity(first_array.len());

    for idx in 0..first_array.len() {
        // 收集所有数组对应位置的元素
        let elements: Vec<&NestedNum> = arrays.iter().map(|arr| &arr[idx]).collect();
        
        match &first_array[idx] {
            // 嵌套数组则递归处理
            NestedNum::Array(_) => {
                let nested_arrays: Vec<Vec<NestedNum>> = elements
                    .iter()
                    .map(|elem| match elem {
                        NestedNum::Array(arr) => arr.clone(),
                        _ => unreachable!("所有数组结构必须一致"),
                    })
                    .collect();
                result.push(NestedNum::Array(add(&nested_arrays)));
            }
            // 数字则直接求和
            NestedNum::Number(_) => {
                let sum: i32 = elements
                    .iter()
                    .map(|elem| match elem {
                        NestedNum::Number(n) => *n,
                        _ => unreachable!("所有数组结构必须一致"),
                    })
                    .sum();
                result.push(NestedNum::Number(sum));
            }
        }
    }

    result
}

fn main() {
    // 测试用例1: [[1]] + [[2]] = [[3]]
    let arr1 = vec![NestedNum::Array(vec![NestedNum::Number(1)])];
    let arr2 = vec![NestedNum::Array(vec![NestedNum::Number(2)])];
    assert_eq!(add(&[arr1, arr2]), vec![NestedNum::Array(vec![NestedNum::Number(3)])]);

    // 测试用例2: [1] + [2] = [3]
    let arr3 = vec![NestedNum::Number(1)];
    let arr4 = vec![NestedNum::Number(2)];
    assert_eq!(add(&[arr3, arr4]), vec![NestedNum::Number(3)]);

    // 测试用例3: [1,2] + [3,4] = [4,6]
    let arr5 = vec![NestedNum::Number(1), NestedNum::Number(2)];
    let arr6 = vec![NestedNum::Number(3), NestedNum::Number(4)];
    assert_eq!(add(&[arr5, arr6]), vec![NestedNum::Number(4), NestedNum::Number(6)]);

    // 测试用例4: [[1,2]] + [[3,4]] = [[4,6]]
    let arr7 = vec![NestedNum::Array(vec![NestedNum::Number(1), NestedNum::Number(2)])];
    let arr8 = vec![NestedNum::Array(vec![NestedNum::Number(3), NestedNum::Number(4)])];
    assert_eq!(add(&[arr7, arr8]), vec![NestedNum::Array(vec![NestedNum::Number(4), NestedNum::Number(6)])]);

    println!("所有测试用例通过!");
}

说明 #

• NestedNum枚举用于统一表示数字和任意深度嵌套数组，保证类型安全。
• add函数要求所有输入数组结构完全一致（对应位置元素类型相同），递归处理嵌套层级，对叶子节点的数字直接求和。

内容的提问来源于stack exchange，提问作者Olgierd Kowalewski