Создание колонны в пандах с использованием векторизации

Question

Я пытаюсь использовать векторизацию на фрейме данных pandas для создания нового столбца. Фрейм данных довольно большой (миллионы записей). Я показываю фиктивный пример здесь. Я показываю не vecotorized версию, которая работает, но не очень эффективна. Я пытаюсь реализовать векторизованную версию при использовании функции (фактическая функция довольно сложна, чем показанная здесь).

import pandas as pd
import numpy as np

df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
df['color'] = np.where(df['Set']=='Z', 'green', 'red')


def test(row):
    if row['color'] =='green':
        value='Green'
    elif row['color'] =='red':
        value=row['Type']
    else: 
        value=row['Set']
    return value

def test1(s,t,c):
    if c =='green':
        value='Green'
    elif c =='red':
        value=t
    else: 
        value=s
    return value

df['new_color']=df.apply(test,axis=1)
#df['new_color']=test1(df.Set,df.Type,df.color)
print(df)

   Set Type color  new_color
0   Z    A  green     Green
1   Z    B  green     Green
2   X    B    red         B
3   Y    C    red         C

Любая помощь будет оценена.

Answer 1

Вы можете сделать с np.where

df['NC']=np.where(df.color=='green','Green',df.Type)

df
Out[1234]: 
  Set Type  color     NC
0   Z    A  green  Green
1   Z    B  green  Green
2   X    B    red      B
3   Y    C    red      C