python - 使用列表过滤 Pandas 数据框的最快方法

Question

假设我有一个 DataFrame，例如:

   col1  col2
0     1     A
1     2     B
2     6     A
3     5     C
4     9     C
5     3     A
6     5     B

以及多个列表，例如:

list_1 = [1, 2, 4]
list_2 = [3, 8]
list_3 = [5, 6, 7, 9]

我可以更新 col2 的值取决于 col1的值是否为包含在列表中，例如:

for i in list_1:
    df.loc[df.col1 == i, 'col2'] = 'A'

for i in list_2:
    df.loc[df.col1 == i, 'col2'] = 'B'

for i in list_3:
    df.loc[df.col1 == i, 'col2'] = 'C'

然而，这是非常缓慢的。对于 30,000 行的数据框，每个列表包含大约 5,000-10,000 个项目，计算可能需要很长时间，尤其是与其他 Pandas 操作相比。有没有更好(更快)的方法来做到这一点？

Answer 1

您可以使用 isin与 np.select这里:

df['col2'] = (np.select([df['col1'].isin(list_1),
                         df['col1'].isin(list_2),
                         df['col1'].isin(list_3)]
                    ,['A','B','C']))

与 Map :

d = dict(zip(map(tuple,[list_1,list_2,list_3]),['A','B','C']))
df['col2'] = df['col1'].map({val: v for k,v in d.items() for val in k})

   col1 col2
0     1    A
1     2    A
2     6    C
3     5    C
4     9    C
5     3    B
6     5    C