在Python Pandas数据框中不迭代实现if else

Question

我想在一个数据框（df）里添加一列。这个新列的值会根据其他列的值来决定。例如：

dc = {'A':[0,9,4,5],'B':[6,0,10,12],'C':[1,3,15,18]}
df = pd.DataFrame(dc)
   A   B   C
0  0   6   1
1  9   0   3
2  4  10  15
3  5  12  18

现在我想添加另一列D，它的值会依赖于A、B和C的值。比如说，如果我在遍历这个数据框，我可以这样做：

for row in df.iterrows():
    if(row['A'] != 0 and row[B] !=0):
         row['D'] = (float(row['A'])/float(row['B']))*row['C']
    elif(row['C'] ==0 and row['A'] != 0 and row[B] ==0):
         row['D'] == 250.0
    else:
         row['D'] == 20.0 

有没有办法可以不使用for循环，或者不使用where()或apply()这些函数来实现呢？

谢谢

Answer 1

这是一个开始：

df['D'] = np.nan
df['D'].loc[df[(df.A != 0) & (df.B != 0)].index] = df.A / df.B.astype(np.float) * df.C

补充一下，如果你没有特别需要用整数的原因，建议你直接把整个东西都转换成浮点数（float）。

df = df.astype(np.float)

这样的话，你就不需要在每次调用的时候都去不停地转换了。

Answer 2

.where 的速度通常比 .apply 快很多，所以如果你只是做一些简单的条件判断（if/else），我建议使用 .where。因为在某些情况下你返回的是单个值（标量），所以使用 np.where 会比 pandas 自带的 .where 更简单。

import pandas as pd
import numpy as np
df['D'] = np.where((df.A!=0) & (df.B!=0), ((df.A/df.B)*df.C),
          np.where((df.C==0) & (df.A!=0) & (df.B==0), 250,
          20))

   A   B   C     D
0  0   6   1  20.0
1  9   0   3  20.0
2  4  10  15   6.0
3  5  12  18   7.5

对于像这样的一个小数据框（df），你不需要担心速度问题。不过，如果你处理的是一个有 10000 行随机数的数据框，那么使用 .where 的速度几乎比上面的 .apply 快 2000 倍：3毫秒对比 5850毫秒。也就是说，如果速度不是问题的话，.apply 通常会更容易理解。

Answer 3

apply这个函数应该能很好地满足你的需求：

In [20]: def func(row):
            if (row == 0).all():
                return 250.0
            elif (row[['A', 'B']] != 0).all():
                return (float(row['A']) / row['B'] ) * row['C']
            else:
                return 20
       ....:     


In [21]: df['D'] = df.apply(func, axis=1)

In [22]: df
Out[22]: 
   A   B   C     D
0  0   6   1  20.0
1  9   0   3  20.0
2  4  10  15   6.0
3  5  12  18   7.5

[4 rows x 4 columns]