python - 函数获取 NumPy 数组中的列数，如果它是一维数组则返回 1

Question

我已经在 3xN NumPy 数组上定义了操作，我想遍历数组的每一列。

我试过:

for i in range(nparray.shape[1]):

但是，如果 nparray.ndim == 1，则失败。

是否有一种干净的方法来确定 NumPy 数组的列数，例如，如果它是 1D 则获取 1 > 数组(就像 MATLAB 的 size 操作一样)？

否则，我已经实现了:

if nparray.ndim == 1:
    num_points = 1
else:
    num_points = nparray.shape[1]

for i in range(num_points):

Answer 1

如果您只是在寻找不那么冗长的内容，您可以这样做:

num_points = np.atleast_2d(nparray).shape[1]

当然，这会创建一个新的临时数组来适应它的形状，这有点愚蠢……但它会非常便宜，因为它只是同一内存的 View 。

但是，我认为您的显式代码更具可读性，除了我可能会使用 try:

try:
    num_points = nparray.shape[1]
except IndexError:
    num_points = 1

如果您重复执行此操作，无论您做什么，都应该将其包装在一个函数中。例如:

def num_points(arr, axis):
    try:
        return arr.shape[axis]
    except IndexError:
        return 1

那么你所要做的就是:

for i in range(num_points(nparray, 1)):

当然，这意味着您只需编辑一个地方就可以在任何地方进行更改，例如:

def num_points(arr, axis):
    return nparray[:,...,np.newaxis].shape[1]