用列表对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]

我可以根据是否将 col1 的值包含在列表中来更新 col2 的值，例如：

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操作相比较的话。有没有更好（更快）的办法来做到这一点？

- Alan

---

https://pandas.pydata.org/docs/user_guide/indexing.html - wwii

Answer 1

您可以在此处使用np.select与isin：

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

- anky

Answer 2

你可以先将列表转换为字典，然后映射到col1。

d1 = {k:'A' for k in list_1}
d2 = {k:'B' for k in list_2}
d3 = {k:'C' for k in list_3}

df['col2'] = (
    df.col1.apply(lambda x: d1.get(x,x))
    .combine_first(df.col1.apply(lambda x: d2.get(x,x)))
    .combine_first(df.col1.apply(lambda x: d2.get(x,x)))
)

如果这些列表中没有重复的元素，你可以将它们合并成一个字典以提高速度：

d = {**{k:'A' for k in list_1}, 
     **{k:'B' for k in list_2}, 
     **{k:'C' for k in list_3}}
df['col2'] = df.col1.apply(lambda x: d.get(x,x))

- Allen Qin

---

我一开始使用了map，但是如果col1的值在字典中不存在，map会返回nan，使用get可以确保值保持不变。 - Allen Qin

很棒的回答 :) +1 - Umar.H

谢谢，比我现有的方法快了3,400倍 :) - Alan

Answer 3

我建议使用字典和条件更新来遍历您的列表：

# Create your update dictionary
col_dict = {
    "A":[1, 2, 4],
    "B":[3, 8],
    "C":[5, 6, 7, 9]
}

# Iterate and update
for key, value in col_dict.items():
  # key is the col name; value is the lookup list
  df["col2"] = np.where(df["col1"].isin(value), key, df["col2"])

存在覆盖值的担忧 - 因为一行数据可以技术上匹配多个列表。如何协调这些更新并不明显。

如果行数据不匹配多个键，则考虑采用动态规划方法，在每次迭代中使用“未匹配”行的运行索引进行更新，随着迭代的进行，更新行数逐渐减少。

- Yaakov Bressler

---

谢谢，比我现有的方法快了3,400倍 :) - Alan

太棒了！@Alan - 整个操作需要多少秒？只是出于好奇。 - Yaakov Bressler

之前，在一个包含2700行数据集上，使用我的for循环方法计算需要3.471秒的时间。但是采用这种方法，只需要0.0009923秒的时间。这是一个巨大的差异，特别是当数据集超过几千行时。以上其他答案的速度都相似，因此我想他们都采用了相同的基本原理。再次感谢 :) - Alan