使用列表过滤 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