我希望能够使用掩码从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循环可以工作但速度很慢。
答案 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]])