NumPy中np.mean、np.max等降维操作如何正确选择axis参数

核心本质：axis是你要「折叠消除」的维度 #

不用死记axis=0沿行、axis=1沿列，记住这个本质逻辑就不会错：你指定的axis值，就是计算后会被消除的维度，剩下的维度顺序和原数组保持一致。
举个最简单的例子：二维数组shape为(样本数, 特征数) = (10, 5)

• 求np.mean(arr, axis=0)：消除第0维（样本维度），输出shape为(5,)，对应每个特征在所有样本上的均值
• 求np.mean(arr, axis=1)：消除第1维（特征维度），输出shape为(10,)，对应每个样本所有特征的均值

结合你的代码逐段解析 #

先明确基础数组的shape：

• 鸢尾花数据集X取前2特征，shape为(150, 2)（150个样本，每个2个特征）
• kneighbors返回的distances、indices shape均为(150, 3)（每个样本对应3个最近邻的距离、索引）

代码段1：平方距离计算 #

dist = np.sum((np.repeat(X, k=3, axis=1).reshape(-1, k, 2) - X[indices]) ** 2, axis=2)

reshape后和X[indices]做差得到的数组shape为(150, 3, 2)，三个维度分别对应「样本、邻居、特征」。
这里指定axis=2是消除特征维度，把每个样本和对应邻居的2个特征的平方差加总，得到欧氏距离的平方，输出shape为(150, 3)，完全符合距离计算的需求。
如果这里选错axis=1，会消除邻居维度，输出(150,2)，得到的是每个样本对3个邻居的特征差总和，完全偏离需求

代码段2：邻近距离最大值 #

k_dist = distances[indices].max(axis=2)

distances[indices]的shape为(150, 3, 3)，指定axis=2消除最后一维，得到每个样本对应邻居的距离最大值，输出shape为(150, 3)。

代码段3：平均邻近距离 #

dist_ = np.mean(distances, axis=1)

distances shape为(150, 3)，指定axis=1消除邻居维度，得到每个样本到3个最近邻的平均距离，输出shape为(150,)，是每个样本对应的一个统计值。
如果这里选错axis=0，会消除样本维度，输出(3,)，得到的是所有样本的第1/2/3近邻的距离均值，和「每个样本的平均邻距」需求完全不符

axis选择的验证方法 #

• 先明确输入数组每个维度的业务含义，不要只记行列对应关系
• 先预判输出的预期shape，反推需要消除的维度下标，就是你要选的axis值
• 拿小尺寸测试数组验证：比如构造一个shape(2,3,2)的小数组，手动计算不同axis的结果，和预期对比即可快速验证

选错axis的影响 #

最直接的影响是输出shape不符合后续代码要求，直接触发维度不匹配报错；就算没有报错，统计维度完全错误，得到的结果也没有业务意义。

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