如何兼容空/1D/2D的NumPy nx3数组，提取其三列？

我有一个NumPy数组nd_arr，存储nx3格式的数据（n取值0到10，数组也可能为空）。我需要一种可靠的方法将其列拆分为三个独立数组。我知道可以用循环追加元素或预定义数组大小来实现，但想寻求更通用的方案。我尝试使用vec_1 = nd_arr[:,0]这类方式提取列，这在n≥2（数组为2D）时有效，但n<2（数组为1D或空）时会触发索引错误（无法识别第二维度）。是否存在类似的、适用于所有n≥0情况的方法？

最小可复现示例（MWE） #

import numpy as np

arr1 = np.array([[1,2,3],[1,2,3],[1,2,3]])  
arr2 = np.array([1,2,3])  
arr3 = np.array([])  

# 每行应返回单独的列a、b、c
a,b,c = arr1[:,0], arr1[:,1], arr1[:,2]  
a,b,c = arr2[:,0], arr2[:,1], arr2[:,2] # 无法运行  
a,b,c = arr3[:,0], arr3[:,1], arr3[:,2] # 无法运行  

当前临时解决方案 #

if len(a:= arr.flatten()) >= 3:
    a,b,c = a[0::3], a[1::3], a[2::3]
else:
    a,b,c = [],[],[]

可以先将数组统一转换为nx3的2D格式，再进行列提取，完美兼容所有n≥0的场景：

import numpy as np

def split_nx3_array(nd_arr):
    # 统一转为2D数组：空数组→(0,3)，1D数组→(1,3)，原2D数组保持nx3形状
    arr_2d = np.asarray(nd_arr).reshape(-1, 3)
    # 安全提取三列
    return arr_2d[:, 0], arr_2d[:, 1], arr_2d[:, 2]

# 测试用例
arr1 = np.array([[1,2,3],[1,2,3],[1,2,3]])  
arr2 = np.array([1,2,3])  
arr3 = np.array([])  

a1, b1, c1 = split_nx3_array(arr1)
a2, b2, c2 = split_nx3_array(arr2)
a3, b3, c3 = split_nx3_array(arr3)

print(a1)  # 输出: [1 1 1]
print(a2)  # 输出: [1]
print(a3)  # 输出: []

关键细节 #

• reshape(-1, 3)会自动计算第一维度的长度：元素总数除以3，保证数组始终是nx3的2D结构；
• 转换后使用[:, idx]索引不会再触发维度错误，空数组会返回三个空的NumPy数组；
• 如果需要返回Python列表而非NumPy数组，只需在返回时添加.tolist()：

  return arr_2d[:, 0].tolist(), arr_2d[:, 1].tolist(), arr_2d[:, 2].tolist()
  

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