如何在N×3矩阵中按指定列值提取其他列数据？

现有如下N×3的numpy数组：

import numpy as np
data=np.array([[30,0.109,1],
               [25,0.517,2],
               [22,0.174,1],
               [35,0.812,3],
               [45,0.215,4],
               [40,0.111,4],
               [50,0.095,4]])

执行以下代码提取第3列数据：

A = data[:,2]
print(A)

得到结果：

[1. 2. 1. 3. 4. 4. 4.]

需要实现：当第3列的值为4时，提取对应行的第1列或第2列数据。预期能通过类似B=data[:,1,[2]=4]的方式实现，目标结果（以第2列为例）为[0.215, 0.111, 0.095]（注：原提问中预期结果存在笔误，对应数据行的第2列实际为此值）。该功能用于数据可视化，第3列数字对应细菌类型，需提取同类型细菌的关联数据进行绘图。

可以使用numpy的布尔索引实现需求，具体步骤如下：

1. 创建布尔掩码，筛选出第3列值为4的行：

mask = data[:, 2] == 4

此时mask的值为[False False False False True True True]，标记出所有符合条件的行。

2. 利用掩码提取指定列的数据：

• 提取第2列（索引为1）：

B = data[mask, 1]
print(B)

输出结果：

[0.215 0.111 0.095]

• 提取第1列（索引为0）：

C = data[mask, 0]
print(C)

输出结果：

[45. 40. 50.]

也可以将两步合并为一行代码：

# 提取第2列
B = data[data[:, 2] == 4, 1]
# 提取第1列
C = data[data[:, 2] == 4, 0]

原理说明 #

numpy的布尔索引会根据掩码数组中的True值，筛选出原数组中对应位置的行，再通过列索引定位到需要提取的列，最终得到符合条件的目标数据，完全满足按分类提取关联数据用于可视化的需求。

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