如果行的某些列在预定义列表中，则从数据框中提取信息

Question

我有一个这样的列表：

list1 = ['4361', '1856', '57586', '79017', '972', '974', '1829', '10787', '85477', '57019', '7431', '53616', '26228', '29085', '5217', '5527']

然后我有两列这样的数据框：

print(df['col A'][0:10])
0      6416
1     84665
2        90
3      2624
4      6118
5       375
6       377
7       377
9       351
10      333


print(df['col B'][0:10])
0      2318
1        88
2      2339
3      5371
4      6774
5     23163
6     23647
7     27236
9     10513
10     1600

我想说'仅返回数据框中的行，如果列表中的项目位于数据框的 col A 或 col B 中'。

我可以想象如何迭代地执行此操作，如下所示：

for each_item in list1:
    for i,row in df.iterrows():
         if each_item in row['col A']:
               print(row)
         if each_item in row['col B']:
               print (row)

我只是想知道是否有一种更简洁的方法可以做到这一点，因为列表和数据框都很大。

我在网上看到了这段代码 sn-p，这将返回 df['col A'] 等于一个值或 df['col B'] 等于一个值的行：

print(df[(df["col A"]==1) | (df_train["col A"]==2)]

如果数据在列表中，我只是不确定如何将其转换为提取数据。有人可以告诉我如何将这种想法融入我的代码中，还是人们认为我的原始代码 sn-p（使用 .iterrows()）是最好的方法？

Answer 1

使用isin:

print(df[(df['col A'].isin(list1)) | (df['col B'].isin(list1))])