如何遍历含子数组的N维numpy数组并生成单值张量集合

核心思路 #

先定位原数组中所有包含子数组的位置，生成这些子数组取值的笛卡尔积，再逐个替换生成仅含单值的N维数组。

代码实现 #

import numpy as np
from itertools import product

def generate_single_value_arrays(original_arr):
    # 收集所有子数组的坐标和对应取值列表
    subarrays_info = []
    for coord, elem in np.ndenumerate(original_arr):
        if isinstance(elem, np.ndarray):
            subarrays_info.append((coord, elem.tolist()))
    
    if not subarrays_info:
        return [original_arr.copy()]
    
    # 生成所有取值组合的笛卡尔积
    values_lists = [vals for _, vals in subarrays_info]
    all_combinations = product(*values_lists)
    
    # 生成每个组合对应的目标数组
    result_arrays = []
    base_arr = np.zeros(original_arr.shape, dtype=np.float64)
    
    for combo in all_combinations:
        current_arr = base_arr.copy()
        for (coord, _), val in zip(subarrays_info, combo):
            current_arr[coord] = val
        result_arrays.append(current_arr)
    
    return result_arrays

使用示例 #

# 构造题目中的输入数组
arr = np.array([
    [np.array([1.  , 3.25, 5.5 , 7.75, 10.  ]), 0, 0],
    [0, np.array([1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]), 0],
    [0, 0, 0]
], dtype=object)

# 生成所有结果数组
result_list = generate_single_value_arrays(arr)

# 打印部分结果示例
print("第一个结果：")
print(result_list[0].astype(int))
print("\n第二个结果：")
print(result_list[1].astype(int))
print("\n最后一个结果：")
print(result_list[-1].astype(int))

关键说明 #

• 兼容任意维度的输入数组，无需预先知晓维度或子数组数量。
• 自动处理不同长度的子数组，生成所有可能的取值组合（数量为各子数组长度的乘积）。
• 返回的每个数组都是独立的numpy对象，可直接传入其他函数使用。
• 若原数组无任何子数组，将返回原数组的副本，避免修改原始数据。

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