Question

我遇到了Python / Pandas的性能问题。我有一个for循环，用于比较Pandas DataFrame中的后续行：

for i in range(1, N):
    if df.column_A.iloc[i] == df.column_A.iloc[i-1]:
        if df.column_B.iloc[i] == 'START' and df.column_B.iloc[i-1] == 'STOP':
            df.time.iloc[i] = df.time.iloc[i] - df.time.iloc[i-1]

哪个可以正常工作，但是速度非常慢。我的数据框有大约100万行，我想知道是否有某种方法可以提高性能。我已经读过关于向量化的知识，但是我不知道从哪里开始。

Answer 1

我认为您可以使用shift和mask：

mask = ((df.column_A == df.column_A.shift()) 
        & (df.column_B == 'START') & (df.column_B.shift() == 'STOP'))
df.loc[mask, 'time'] -= df.time.shift().loc[mask]

掩码选择“ column_A”中的值等于上一个值（由shift获得），“ column_B”等于“ START”且上一行为“ STOP”的行'。使用loc可让您在“时间”列中通过mask更改所有选定行的值，方法是使用相同的掩码删除前一行（再次为shift）的值。列时间

编辑：例如：

df = pd.DataFrame({'column_A': [0,1,1,2,1,2,2], 'column_B': ['START', 'STOP', 'START','STOP', 'START','STOP', 'START'], 'time':range(7)})
   column_A column_B  time
0         0    START     0
1         1     STOP     1
2         1    START     2
3         2     STOP     3
4         1    START     4
5         2     STOP     5
6         2    START     6

因此，第2行和第6行符合您的条件，因为前一行在column_A中具有相同的值，而在column_B中具有“ START”，而先前的行具有“ STOP”。

运行代码后，您将获得df：

   column_A column_B  time
0         0    START   0.0
1         1     STOP   1.0
2         1    START   1.0
3         2     STOP   3.0
4         1    START   4.0
5         2     STOP   5.0
6         2    START   1.0

第2行的时间值为1（最初是第1行的2减去值），而第6行的时间值为6（6-5）

编辑以进行时间比较，让我们创建具有3000行的df

df = pd.DataFrame( [['A', 'START', 3], ['A', 'STOP', 6], ['B', 'STOP', 2], 
                    ['C', 'STOP', 1], ['C', 'START', 9], ['C', 'STOP', 7]],
                   columns=['column_A', 'column_B', 'time'] )
df = pd.concat([df]*500)
df.shape
Out[16]: (3000, 3)

现在使用两种方法创建两个函数：

# original method
def espogian (df):
    N = df.shape[0]
    for i in range(1, N):
        if df.column_A.iloc[i] == df.column_A.iloc[i-1]:
            if df.column_B.iloc[i] == 'START' and df.column_B.iloc[i-1] == 'STOP':
                df.time.iloc[i] = df.time.iloc[i] - df.time.iloc[i-1]
    return df
# mine
def ben(df):
    mask = ((df.column_A == df.column_A.shift()) 
        & (df.column_B == 'START') & (df.column_B.shift() == 'STOP'))
    df.loc[mask, 'time'] -= df.time.shift().loc[mask]
    return df

并运行timeit：

%timeit espogian (df)
1 loop, best of 3: 8.71 s per loop

%timeit ben (df)
100 loops, best of 3: 4.79 ms per loop

# verify they are equal
df1 = espogian (df)
df2 = ben (df)
(df1==df2).all()
Out[24]: 
column_A    True
column_B    True
time        True

提高for循环比较熊猫数据框行的性能

1 个答案: