通過列值（Python pandas）在DataFrame切片上查找計算函數的最快方法

Question

我試圖在數據框上創建一個列，其中包含列A（值列）的最小值，列B（id列）具有特定值。 我的代碼很慢。 我正在尋找一種更快的方法來做到這一點。 這是我的小功能：

def apply_by_id_value(df, id_col="id_col", val_col="val_col", offset_col="offset", f=min):
    for rid in set(df[id_col].values):
        df.loc[df[id_col] == rid, offset_col] =  f(df[df[id_col] == rid][val_col])
    return df

示例用法：

import pandas as pd
import numpy as np
# create data frame
df = pd.DataFrame({"id_col":[0, 0, 0, 1, 1, 1, 2, 2, 2], 
                   "val_col":[0.1, 0.2, 0.3, 0.6, 0.4, 0.5, 0.2, 0.1, 0.0]})

print df.head(10)
# output
   id_col  val_col
0       0      0.1
1       0      0.2
2       0      0.3
3       1      0.6
4       1      0.4
5       1      0.5
6       2      0.2
7       2      0.1
8       2      0.0

df = apply_by_id_value(df)
print df.head(10) 
# output

   id_col  val_col  offset
0       0      0.1     0.1
1       0      0.2     0.1
2       0      0.3     0.1
3       1      0.6     0.4
4       1      0.4     0.4
5       1      0.5     0.4
6       2      0.2     0.0
7       2      0.1     0.0
8       2      0.0     0.0

更多上下文：在我的實際數據中，“id_col”列有大約30000個或更多的唯一值。 這意味着數據幀必須切片30000次。 我想這是瓶頸。

Answer 1

在'id_col'上執行groupby ，然后執行transform傳遞函數'min'，這將返回與原始df對齊的Series，以便您可以添加為新列：

In [13]:

df = pd.DataFrame({"id_col":[0, 0, 0, 1, 1, 1, 2, 2, 2], 
                   "val_col":[0.1, 0.2, 0.3, 0.6, 0.4, 0.5, 0.2, 0.1, 0.0]})
df['offset'] = df.groupby('id_col').transform('min')
df
Out[13]:
   id_col  val_col  offset
0       0      0.1     0.1
1       0      0.2     0.1
2       0      0.3     0.1
3       1      0.6     0.4
4       1      0.4     0.4
5       1      0.5     0.4
6       2      0.2     0.0
7       2      0.1     0.0
8       2      0.0     0.0

計時

In [15]:

def apply_by_id_value(df, id_col="id_col", val_col="val_col", offset_col="offset", f=min):
    for rid in set(df[id_col].values):
        df.loc[df[id_col] == rid, offset_col] =  f(df[df[id_col] == rid][val_col])
    return df
%timeit apply_by_id_value(df)
%timeit df.groupby('id_col').transform('min')
100 loops, best of 3: 8.12 ms per loop
100 loops, best of 3: 5.99 ms per loop

因此groupby和transform在這個數據集上更快，我希望它在你的真實數據集上明顯更快，因為它會更好地擴展。

對於800,000行df，我得到以下時間：

1 loops, best of 3: 611 ms per loop
1 loops, best of 3: 438 ms per loop