如何在Python的if语句中判断数组维度为(p,)？

在NumPy中，存在两种容易混淆的数组维度形式：(p,)（真正的一维数组）和(p, 1)（二维数组的单列形式）。直接通过shape[1]访问第二个维度时，(p,)类型的数组会触发索引错误，以下是几种无错误的判断方法：

方法1：检查数组的维度数（ndim） #

(p,)类型的数组维度数为1，(p,1)类型的为2，直接通过ndim属性判断：

import numpy as np

x = np.array([1, 2, 3])
print(x.ndim)  # 输出1，对应(p,)类型

x2 = np.random.rand(3, 1)
print(x2.ndim) # 输出2，对应(p,1)类型

# 判断逻辑
def is_1d(arr):
    return arr.ndim == 1

print(is_1d(x))  # True
print(is_1d(x2)) # False

方法2：检查shape元组的长度 #

(p,)的shape是长度为1的元组，(p,1)的shape是长度为2的元组，通过len(shape)判断：

# 判断逻辑
def is_1d(arr):
    return len(arr.shape) == 1

print(is_1d(x))  # True
print(is_1d(x2)) # False

额外提示 #

如果需要统一两种维度的数组，可以使用np.squeeze()将(p,1)转换为(p,)，或者用arr.reshape(-1,1)将(p,)转换为(p,1)，避免后续维度操作出错。

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