python - 与列表的 Pandas DataFrame 列的比较

Question

我有一个像这样的数据框 df:

col1 | col2
 a   | [1,2]
 b   | [3,4]
 c   | [3,9]

我想根据匹配的输入数组获取行，所以如果我有数组 [1,2]，我可以获得:

col1 | col2
 a   | [1,2]

当我尝试使用此公式执行此操作时，它不起作用:

df.loc[df['Col2'] == [1,2]]
Error: Lengths must match to compare

Answer 1

错误的真正原因是并非所有列表的大小都相同，这会导致 DataFrame.eq 出现问题。

解决这个问题的最佳方法是使用列表理解构建一个 bool 掩码，然后使用它索引到 df:

df[[v == [1, 2] for v in df['col2'].tolist()]]

另一种选择是 df.apply，但它的速度远不如这个。