如何统一获取1维或2维numpy数组的行数

我们默认约定1维数组视为1行结构，2维数组取第一维长度作为行数，可通过以下两种常用方法实现统一计算：

• 方法1：使用np.atleast_2d()统一转换维度后计算
  该方法会自动将输入数组转为至少2维的格式，1维数组会被调整为shape为(1, n)的2维数组，无需手动分支判断，代码最简洁：

import numpy as np

# 1维数组测试
x3 = np.array([1, 2, 3, 3, 4, 5])
row_num = np.atleast_2d(x3).shape[0]
print(row_num) # 输出1

# 2维数组测试
x4 = np.array([[1, 2, 3], [3, 4, 5]])
row_num = np.atleast_2d(x4).shape[0]
print(row_num) # 输出2

• 方法2：通过ndim属性判断维度后分支计算
  逻辑清晰直观，适合需要自定义不同维度行数规则的场景：

def get_array_rows(arr):
    # 1维数组返回行数为1
    if arr.ndim == 1:
        return 1
    # 2维及以上数组返回第一维长度
    return arr.shape[0]

print(get_array_rows(x3)) # 输出1
print(get_array_rows(x4)) # 输出2

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