pandas:根据掩码有条件地为每列选择一个行单元格

时间:2017-04-18 19:32:33

标签: pandas numpy conditional mask

我希望能够使用掩码从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],我想返回:

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

但是以pythonic的方式。

我试过了例如:

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

给出了我不理解的Type错误。

for循环可以工作但速度很慢。

1 个答案:

答案 0 :(得分:2)

使用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]])