pandas: 有条件地 select 基于掩码的每一列的行单元格

Question

我希望能够使用掩码从 pandas 数据帧中提取值。但是，经过四处寻找，我找不到解决问题的方法。

df = pd.DataFrame(np.random.randint(0,2, size=(2,10)))
mask = np.random.randint(0,2, size=(1,10))

我基本上希望掩码用作每一列的索引查找。

所以如果 [a,b] 列的掩码是 [0,1]，我想 return:

df.iloc[0,a], df.iloc[1,b]

但以 pythonic 方式。

我试过例如：

df.apply(lambda x: df.iloc[mask[x], x] for x in range(len(mask)))

它给出了一个我不理解的类型错误。

for 循环可以工作但是很慢。

Answer 1

使用 NumPy，这被称为 advanced-indexing 并且应该非常有效 -

df.values[mask, np.arange(mask.size)]

样本运行-

In [59]: df = pd.DataFrame(np.random.randint(11,99, size=(5,10)))

In [60]: mask = np.random.randint(0,5, size=(1,10))

In [61]: df
Out[61]: 
    0   1   2   3   4   5   6   7   8   9
0  17  87  73  98  32  37  61  58  35  87
1  52  64  17  79  20  19  89  88  19  24
2  50  33  41  75  19  77  15  59  84  86
3  69  13  88  78  46  76  33  79  27  22
4  80  64  17  95  49  16  87  82  60  19

In [62]: mask
Out[62]: array([[2, 3, 0, 4, 2, 2, 4, 0, 0, 0]])

In [63]: df.values[mask, np.arange(mask.size)]
Out[63]: array([[50, 13, 73, 95, 19, 77, 87, 58, 35, 87]])

pandas: 有条件地 select 基于掩码的每一列的行单元格

pandas: conditionally select a row cell for each column based on a mask

conditional

numpy

mask

pandas