如何按照Python规则在pandas DataFrame中新增列并修复赋值为NaN的问题

现有数据 #

构造DataFrame的代码如下：

import numpy as np
import pandas as pd

df = pd.DataFrame({'Name': ['Station', 'Station', 'Sensor', 'Sensor', 
                            'Sensor', 'Sensor', 'Station', 'Station'], 
                  'id': [10, 10, 11, 11, 12, 12, 13, 13]})

该DataFrame输出内容：

Name  id
0  Station  10
1  Station  10
2   Sensor  11
3   Sensor  11
4   Sensor  12
5   Sensor  12
6  Station  13
7  Station  13

同时有id参考数组：

array_id = np.array([10, 11])

需求说明 #

为DataFrame新增名为class的列，赋值规则为：

• 若该行id属于array_id，赋值为class1
• 若该行id不属于array_id，赋值为class2

原代码问题 #

原有嵌套循环写法存在两个核心问题：

1. 逻辑错误：内层循环遍历array_id时，只要最后一个对比的元素和当前行id不匹配，就算前面有匹配项，结果也会被覆盖为class2
2. 赋值失效：df['class'].iloc[i]属于链式索引操作，pandas会返回数据副本而非原对象视图，赋值操作不会同步到原DataFrame，最终导致class列全为NaN

解决方案 #

直接用pandas+numpy的矢量化操作实现，仅需一行代码，执行效率远高于循环写法：

df['class'] = np.where(df['id'].isin(array_id), 'class1', 'class2')

执行后输出结果符合预期：

Name  id   class
0  Station  10  class1
1  Station  10  class1
2   Sensor  11  class1
3   Sensor  11  class1
4   Sensor  12  class2
5   Sensor  12  class2
6  Station  13  class2
7  Station  13  class2

内容的提问来源于stack exchange，提问作者Jane Borges