我如何根据某列的条件来赋值？

Question

我在学习dataframe的时候碰到一个基本的问题，

如果data列>0，我需要把a列的相同行的值赋给data2列，如果data2列<0，我需要把b列的相同行的值赋给data2，我尝试了好几种办法，都没成功

代码如下(python 3.8)：

import pandas as pd
import numpy as np

data = {"data": np.random.randn(200)}
df = pd.DataFrame(data)
df['a']=np.random.randint(1,100, len(df))
df['b']=np.random.randint(1,100, len(df))
df['data2']=0

上面的df和我的数据类似，我想到一个方法，感觉很粗糙，有没有更优的方法？

df['data2'].loc[df['data']>0]=df[df['data']>0]['a']
df['data2'].loc[df['data']<0]=df[df['data']<0]['b']

Answer 1

np.where最简单：

import pandas as pd
import numpy as np


def main():
    condition = {"condition": np.random.randn(200)}
    df = pd.DataFrame(condition)
    df['data1'] = np.random.randint(1, 100, len(df))
    df['data2'] = np.random.randint(1, 100, len(df))
    df['result'] = np.where(df['condition'] > 0, df['data1'], df['data2'])
    print(df.head(10))
    return


main()