Pandas有条件地创建一个序列/数据帧列

我如何添加一个color列到下面的数据帧,以便color='green'如果Set == 'Z',和color='red'否则?

    Type       Set
1    A          Z
2    B          Z
3    B          X
4    C          Y
743310 次浏览

如果你只有两种选择:

df['color'] = np.where(df['Set']=='Z', 'green', 'red')

例如,

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')
print(df)

收益率

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

如果你有两个以上的条件,那么使用np.select。例如,如果你想要color

  • (df['Set'] == 'Z') & (df['Type'] == 'A')yellow
  • 否则当(df['Set'] == 'Z') & (df['Type'] == 'B')blue
  • 否则当(df['Type'] == 'B')purple
  • 否则black,

然后使用

df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
conditions = [
(df['Set'] == 'Z') & (df['Type'] == 'A'),
(df['Set'] == 'Z') & (df['Type'] == 'B'),
(df['Type'] == 'B')]
choices = ['yellow', 'blue', 'purple']
df['color'] = np.select(conditions, choices, default='black')
print(df)

的收益率

  Set Type   color
0   Z    A  yellow
1   Z    B    blue
2   X    B  purple
3   Y    C   black

另一种实现这一目标的方法是

df['color'] = df.Set.map( lambda x: 'red' if x == 'Z' else 'green')

列表推导式是有条件地创建另一列的另一种方法。如果您在列中使用对象dtype,就像您的示例一样,列表推导式通常优于大多数其他方法。

示例列表理解:

df['color'] = ['red' if x == 'Z' else 'green' for x in df['Set']]

%时间测试:

import pandas as pd
import numpy as np


df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
%timeit df['color'] = ['red' if x == 'Z' else 'green' for x in df['Set']]
%timeit df['color'] = np.where(df['Set']=='Z', 'green', 'red')
%timeit df['color'] = df.Set.map( lambda x: 'red' if x == 'Z' else 'green')


1000 loops, best of 3: 239 µs per loop
1000 loops, best of 3: 523 µs per loop
1000 loops, best of 3: 263 µs per loop

下面的方法比计时在这里的方法慢,但我们可以基于多个列的内容计算额外的列,并且可以为额外的列计算两个以上的值。

使用“Set”列的简单示例:

def set_color(row):
if row["Set"] == "Z":
return "red"
else:
return "green"


df = df.assign(color=df.apply(set_color, axis=1))


print(df)
  Set Type  color
0   Z    A    red
1   Z    B    red
2   X    B  green
3   Y    C  green

考虑到更多颜色和更多列的例子:

def set_color(row):
if row["Set"] == "Z":
return "red"
elif row["Type"] == "C":
return "blue"
else:
return "green"


df = df.assign(color=df.apply(set_color, axis=1))


print(df)
  Set Type  color
0   Z    A    red
1   Z    B    red
2   X    B  green
3   Y    C   blue

编辑(21/06/2019):使用plydata

也可以使用plydata来做这类事情(不过,这似乎比使用assignapply还要慢)。

from plydata import define, if_else

简单的if_else:

df = define(df, color=if_else('Set=="Z"', '"red"', '"green"'))


print(df)
  Set Type  color
0   Z    A    red
1   Z    B    red
2   X    B  green
3   Y    C  green

嵌套if_else:

df = define(df, color=if_else(
'Set=="Z"',
'"red"',
if_else('Type=="C"', '"green"', '"blue"')))


print(df)
  Set Type  color
0   Z    A    red
1   Z    B    red
2   X    B   blue
3   Y    C  green

这是另一种方法,使用字典将新值映射到列表中的键:

def map_values(row, values_dict):
return values_dict[row]


values_dict = {'A': 1, 'B': 2, 'C': 3, 'D': 4}


df = pd.DataFrame({'INDICATOR': ['A', 'B', 'C', 'D'], 'VALUE': [10, 9, 8, 7]})


df['NEW_VALUE'] = df['INDICATOR'].apply(map_values, args = (values_dict,))

它看起来像什么:

df
Out[2]:
INDICATOR  VALUE  NEW_VALUE
0         A     10          1
1         B      9          2
2         C      8          3
3         D      7          4

当你有很多__abc0类型的语句要做(即许多唯一的值要替换)时,这种方法非常强大。

当然你可以这样做:

df['NEW_VALUE'] = df['INDICATOR'].map(values_dict)

但是在我的机器上,这种方法比上面的apply方法慢三倍多。

你也可以这样做,使用dict.get:

df['NEW_VALUE'] = [values_dict.get(v, None) for v in df['INDICATOR']]

你可以简单地使用强大的.loc方法,并根据需要使用一个或多个条件(用pandas=1.0.5测试)。

代码总结:

df=pd.DataFrame(dict(Type='A B B C'.split(), Set='Z Z X Y'.split()))
df['Color'] = "red"
df.loc[(df['Set']=="Z"), 'Color'] = "green"


#practice!
df.loc[(df['Set']=="Z")&(df['Type']=="B")|(df['Type']=="C"), 'Color'] = "purple"


解释:

df=pd.DataFrame(dict(Type='A B B C'.split(), Set='Z Z X Y'.split()))


# df so far:
Type Set
0    A   Z
1    B   Z
2    B   X
3    C   Y

添加“color”列,并将所有值设置为“red”;

df['Color'] = "red"

应用你的单一条件:

df.loc[(df['Set']=="Z"), 'Color'] = "green"




# df:
Type Set  Color
0    A   Z  green
1    B   Z  green
2    B   X    red
3    C   Y    red

或者多重条件:

df.loc[(df['Set']=="Z")&(df['Type']=="B")|(df['Type']=="C"), 'Color'] = "purple"

你可以在这里阅读Pandas逻辑运算符和条件选择: Pandas中用于布尔索引的逻辑运算符 < / p >

.apply()方法的一行代码如下:

df['color'] = df['Set'].apply(lambda set_: 'green' if set_=='Z' else 'red')

在此之后,df数据帧看起来像这样:

>>> print(df)
Type Set  color
0    A   Z  green
1    B   Z  green
2    B   X    red
3    C   Y    red

如果你在处理海量数据,记忆方法是最好的:

# First create a dictionary of manually stored values
color_dict = {'Z':'red'}


# Second, build a dictionary of "other" values
color_dict_other = {x:'green' for x in df['Set'].unique() if x not in color_dict.keys()}


# Next, merge the two
color_dict.update(color_dict_other)


# Finally, map it to your column
df['color'] = df['Set'].map(color_dict)

我的一般经验法则是记住什么时候:n_distinct & lt;data_size/4

在一种情况下,记忆10,000行,不同值不超过2,500。

你可以使用pandas方法wheremask:

df['color'] = 'green'
df['color'] = df['color'].where(df['Set']=='Z', other='red')
# Replace values where the condition is False

df['color'] = 'red'
df['color'] = df['color'].mask(df['Set']=='Z', other='green')
# Replace values where the condition is True

或者,你可以使用方法transform和lambda函数:

df['color'] = df['Set'].transform(lambda x: 'green' if x == 'Z' else 'red')

输出:

  Type Set  color
1    A   Z  green
2    B   Z  green
3    B   X    red
4    C   Y    red

@chai的性能比较:

import pandas as pd
import numpy as np
df = pd.DataFrame({'Type':list('ABBC')*1000000, 'Set':list('ZZXY')*1000000})
 

%timeit df['color1'] = 'red'; df['color1'].where(df['Set']=='Z','green')
%timeit df['color2'] = ['red' if x == 'Z' else 'green' for x in df['Set']]
%timeit df['color3'] = np.where(df['Set']=='Z', 'red', 'green')
%timeit df['color4'] = df.Set.map(lambda x: 'red' if x == 'Z' else 'green')


397 ms ± 101 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
976 ms ± 241 ms per loop
673 ms ± 139 ms per loop
796 ms ± 182 ms per loop

如果只有2选择,则使用np.where()

df = pd.DataFrame({'A':range(3)})
df['B'] = np.where(df.A>2, 'yes', 'no')

如果你有超过2选择,也许apply()可以工作 输入< / p >

arr = pd.DataFrame({'A':list('abc'), 'B':range(3), 'C':range(3,6), 'D':range(6, 9)})

arr是

    A   B   C   D
0   a   0   3   6
1   b   1   4   7
2   c   2   5   8

如果你想让列E是if arr.A =='a' then arr.B elif arr.A=='b' then arr.C elif arr.A == 'c' then arr.D else something_else

arr['E'] = arr.apply(lambda x: x['B'] if x['A']=='a' else(x['C'] if x['A']=='b' else(x['D'] if x['A']=='c' else 1234)), axis=1)

最后是arr

    A   B   C   D   E
0   a   0   3   6   0
1   b   1   4   7   4
2   c   2   5   8   8

pyjanitor中的case_when函数是pd.Series.mask的包装器,并为多种条件提供了可链接/方便的形式:

对于单一条件:

df.case_when(
df.col1 == "Z",  # condition
"green",         # value if True
"red",           # value if False
column_name = "color"
)


Type Set  color
1    A   Z  green
2    B   Z  green
3    B   X    red
4    C   Y    red

适用于多种情况:

df.case_when(
df.Set.eq('Z') & df.Type.eq('A'), 'yellow', # condition, result
df.Set.eq('Z') & df.Type.eq('B'), 'blue',   # condition, result
df.Type.eq('B'), 'purple',                  # condition, result
'black',              # default if none of the conditions evaluate to True
column_name = 'color'
)
Type  Set   color
1    A   Z  yellow
2    B   Z    blue
3    B   X  purple
4    C   Y   black

更多的例子可以在在这里中找到

使用np.select的更简洁的方法:

a = np.array([['A','Z'],['B','Z'],['B','X'],['C','Y']])
df = pd.DataFrame(a,columns=['Type','Set'])


conditions = [
df['Set'] == 'Z'
]


outputs = [
'Green'
]
# conditions Z is Green, Red Otherwise.
res = np.select(conditions, outputs, 'Red')
res
array(['Green', 'Green', 'Red', 'Red'], dtype='<U5')
df.insert(2, 'new_column',res)


df
Type    Set new_column
0   A   Z   Green
1   B   Z   Green
2   B   X   Red
3   C   Y   Red


df.to_numpy()
    

array([['A', 'Z', 'Green'],
['B', 'Z', 'Green'],
['B', 'X', 'Red'],
['C', 'Y', 'Red']], dtype=object)


%%timeit conditions = [df['Set'] == 'Z']
outputs = ['Green']
np.select(conditions, outputs, 'Red')


134 µs ± 9.71 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)


df2 = pd.DataFrame({'Type':list('ABBC')*1000000, 'Set':list('ZZXY')*1000000})
%%timeit conditions = [df2['Set'] == 'Z']
outputs = ['Green']
np.select(conditions, outputs, 'Red')


188 ms ± 26.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)