python - Pandas 在数据框中使用二维数组

Question

我有一个 Pandas DataFrame，其中包含一个 2D 数组作为列，如下所示:

Name    2DValueList
item 1  [ [ 0.0, 1.0 ], [ 0.0, 6.0 ], [ 0.0, 2.0 ] ]
item 2  [ [ 0.0, 2.0 ], [ 0.0, 1.0 ], [ 0.0, 1.0 ] ]
item 3  [ [ 0.0, 1.0 ], [ 0.0, 3.0 ], [ 0.0, 5.0 ], [ 0.0, 1.0 ] ]
item 4  
item 5  [ [ 0.0, 4.0 ], [ 0.0, 1.0 ], [ 0.0, 2.0 ] ]

第一个值与这个问题无关，所以我只是将它们全部设为 0。我只对第二个值感兴趣。另请注意，对的数量可能会有所不同或为空。

我希望能够创建一个仅包含数组中前(最大)n 个元素的新数据框。

前 2 个元素看起来像这样:

Name    2DValueList
item 1  [ [ 0.0, 6.0 ], [ 0.0, 2.0 ] ]
item 2  [ [ 0.0, 2.0 ], [ 0.0, 1.0 ] ]
item 3  [ [ 0.0, 5.0 ], [ 0.0, 3.0 ] ]
item 4
item 5  [ [ 0.0, 4.0 ], [ 0.0, 2.0 ] ]

我会使用 pandas nlargest，但我不确定如何让它接受二维数组的列。

实际上，二维数组包含数千个值对，并且有数万行。我愿意接受更好的方法来保存这些更通用的数据。

Answer 1

如果 2DValueList 的每个单元格是列表的列表，有效的方法是使用 heapq.nlargest与 itemgetter与列表理解一起

from heapq import nlargest
from operator import itemgetter

df['new_list'] = [nlargest(2, x, key=itemgetter(1)) for x in df['2DValueList']]

Out[119]:
     Name               2DValueList            new_list
0  item 1  [[0, 1], [0, 6], [0, 2]]    [[0, 6], [0, 2]]
1  item 2  [[0, 2], [0, 1], [0, 1]]    [[0, 2], [0, 1]]
2  item 3  [[0, 1], [0, 3], [0, 5]]    [[0, 5], [0, 3]]
3  item 4  [[0, 4], [0, 1], [0, 2]]    [[0, 4], [0, 2]]

如果每个单元格都是 numpy 2darray，上述方法仍然可以正常工作。但是，我认为使用 numpy argsort会更好

df['new_list'] = [x[np.argsort(-x, axis=0)[:2,1]] for x in df['2DValueList']]

Out[128]:
     Name               2DValueList            new_list
0  item 1  [[0, 1], [0, 6], [0, 2]]    [[0, 6], [0, 2]]
1  item 2  [[0, 2], [0, 1], [0, 1]]    [[0, 2], [0, 1]]
2  item 3  [[0, 1], [0, 3], [0, 5]]    [[0, 5], [0, 3]]
3  item 4  [[0, 4], [0, 1], [0, 2]]    [[0, 4], [0, 2]]

最后，如果不需要按排序顺序排列的前 n 个最大的子数组，argpartition会比 argsort 更快