如何提取Pandas DataFrame条件子集及对应Numpy数组的同索引行

解决方案 #

你直接利用筛选df时生成的布尔掩码就能提取numpy数组的对应条目，完全不需要考虑df的索引是否连续，两种常用实现方式如下：

方法1：共用筛选掩码（最简便）

筛选df时生成的布尔Series和ar的行顺序完全对齐，直接转成numpy布尔数组索引ar即可：

# 生成共用的筛选布尔掩码
mask = df['column'] == value
# 提取df子集
subdf = df[mask]
# 提取ar对应条目，和subdf行顺序完全一致
sub_ar = ar[mask.to_numpy()]

原理说明：布尔掩码mask的长度和df、ar完全相同，每个位置的布尔值对应df同一位置的行是否符合筛选条件，和df的索引值是否是连续整数没有关联。

方法2：基于已有subdf提取

如果你已经提前生成了subdf，可以通过索引定位拿到对应位置：

# 已提前生成subdf的场景
subdf = df[df['column'] == value]
# 获取subdf的行在原df中的位置下标
row_pos = df.index.get_indexer(subdf.index)
# 提取ar对应条目
sub_ar = ar[row_pos]

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