如何处理含子数组与整数的混合NumPy数组：移除整数或转整数为数组？

问题背景 #

现有一个包含子数组和整数的NumPy混合对象数组，示例如下：

import numpy as np
mixed_array = np.array([np.array([1, 2]), np.array([1, 2]), 0])

需求有两个：

1. 高效筛选出数组中的所有子数组，得到仅包含子数组的结果：

   filtered_array = np.array([np.array([1, 2]), np.array([1, 2])])
   

2. 将数组中的整数转换为单元素NumPy数组（如把0转为np.array([0])）。

由于处理的是超大数组，需避免低效循环。

解决方案 #

1. 高效筛选子数组 #

利用布尔索引结合类型判断生成掩码，这是针对object数组的高效处理方式：

# 生成布尔掩码：标记元素是否为NumPy数组
mask = np.array([isinstance(item, np.ndarray) for item in mixed_array])
# 应用掩码筛选子数组
filtered_array = mixed_array[mask]

说明：列表推导式生成掩码的开销远低于显式循环，配合NumPy的布尔索引（矢量化操作），能满足超大数组的性能要求。

2. 将整数转换为单元素数组 #

可以通过列表推导式直接完成转换，性能同样优于逐元素循环：

# 遍历元素，整数转单元素数组，子数组保持不变
converted_array = np.array(
    [item if isinstance(item, np.ndarray) else np.array([item]) 
     for item in mixed_array]
)

如果偏好NumPy风格的写法，也可以用np.vectorize包装转换逻辑：

def convert_element(item):
    return item if isinstance(item, np.ndarray) else np.array([item])

# 矢量化转换函数
vec_convert = np.vectorize(convert_element, otypes=[object])
converted_array = vec_convert(mixed_array)

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