Pandas DataFrame-根据其他列分配1,0个值

Question

我有一个包含国名及其能源输出百分比的数据框。 我需要添加一个新列，根据该国的能源输出是高于还是低于能源输出的中位数来分配1或0。 一些伪代码是：

import pandas as pd
def answer():
    df = pd.DataFrame({'name':['china', 'america', 'canada'], 'output': [33.2, 15.0, 5.0]})
    df['newcol'] = df.where(df['output'] > df['output'].median(), 1, 0)
    return df['newcol']
answer()

代码返回

ValueError：传递的项目数错误2，展示位置暗含1

我觉得这是一个非常简单的修复程序，但是我Pandas 。 请帮助我结束挫败感

Answer 1

您无需循环，因为解决方案是矢量化的。

df['newcol'] = np.where((df['output'] > df['output'].median()), 1, 0)

    name    output  newcol
0   china   33.2    1
1   america 15.0    0
2   canada  5.0     0

对于传递的错误错误项目数，df.where与np.where略有不同。 它返回与self形状相同的对象，其对应的条目来自cond为True的self，否则来自其他对象。 因此，它在您的情况下返回的数据帧是两列而不是序列，因此当您尝试将该数据帧分配给序列时，会收到错误消息。

Answer 2

@Vaishali解释了为什么pd.DataFrame.where无法按您预期的那样工作，并建议您改用np.where ，这是非常好的建议。

我提供的是，您可以将布尔结果简单地转换为整数。

设定

df = pd.DataFrame({
    'name':['china', 'america', 'canada'],
    'output': [33.2, 15.0, 5.0]
})

选项1

df['newcol'] = (df['output'] > df['output'].median()).astype(int)

---

选项2
或者通过使用基础的numpy数组更快

o = df['output'].values
df['newcol'] = (o > np.median(o)).astype(int)