在蟒蛇熊猫群:快速方式

Question

我想改进python pandas中groupby的时间。我有这样的代码:

df["Nbcontrats"] = df.groupby(['Client', 'Month'])['Contrat'].transform(len)

目标是计算一个客户在一个月内有多少合同，并将这些信息添加到新列(Nbcontrats)中。

• Client:客户端代码
• Month:数据提取月份
Contrat:合同号

我想改善时间。下面我只处理真实数据的一个子集:

%timeit df["Nbcontrats"] = df.groupby(['Client', 'Month'])['Contrat'].transform(len)
1 loops, best of 3: 391 ms per loop
df.shape
Out[309]: (7464, 61)

如何提高执行时间?

Answer 1

有一种方法:

• 将输入数据帧中的相关列(['Client', 'Month'])切片到NumPy数组中。这主要是一个以性能为中心的想法，因为我们稍后将使用NumPy函数，这些函数被优化为与NumPy数组一起工作。

• 将['Client', 'Month']中的两列数据转换为单个1D数组，这将是一个线性索引，考虑两列中的元素为对。因此，我们可以假设'Client'中的元素表示行索引，而'Month'中的元素是列索引。这就像从2D到1D。但是，问题将是决定二维网格的形状来执行这样的映射。为了涵盖所有对，一个安全的假设是假设一个2D网格，由于Python中基于0的索引，其每列的维度都比最大值大1。因此，我们将得到线性索引。

• 接下来，我们根据每个线性索引的唯一性对其进行标记。我认为这将对应于grouby获得的密钥。我们还需要沿着1D数组的整个长度获得每个组/唯一键的计数。最后，用这些标记索引计数应该映射每个元素各自的计数。

就是这个意思!下面是实现-

# Save relevant columns as a NumPy array for performing NumPy operations afterwards
arr_slice = df[['Client', 'Month']].values
# Get linear indices equivalent of those columns
lidx = np.ravel_multi_index(arr_slice.T,arr_slice.max(0)+1)
# Get unique IDs corresponding to each linear index (i.e. group) and grouped counts
unq,unqtags,counts = np.unique(lidx,return_inverse=True,return_counts=True)
# Index counts with the unique tags to map across all elements with the counts
df["Nbcontrats"] = counts[unqtags]

运行时测试

1)定义函数:

def original_app(df):
    df["Nbcontrats"] = df.groupby(['Client', 'Month'])['Contrat'].transform(len)
def vectorized_app(df):
    arr_slice = df[['Client', 'Month']].values
    lidx = np.ravel_multi_index(arr_slice.T,arr_slice.max(0)+1)
    unq,unqtags,counts = np.unique(lidx,return_inverse=True,return_counts=True)
    df["Nbcontrats"] = counts[unqtags]

2)验证结果:

In [143]: # Let's create a dataframe with 100 unique IDs and of length 10000
     ...: arr = np.random.randint(0,100,(10000,3))
     ...: df = pd.DataFrame(arr,columns=['Client','Month','Contrat'])
     ...: df1 = df.copy()
     ...: 
     ...: # Run the function on the inputs
     ...: original_app(df)
     ...: vectorized_app(df1)
     ...: 
In [144]: np.allclose(df["Nbcontrats"],df1["Nbcontrats"])
Out[144]: True

3)最后计时:

In [145]: # Let's create a dataframe with 100 unique IDs and of length 10000
     ...: arr = np.random.randint(0,100,(10000,3))
     ...: df = pd.DataFrame(arr,columns=['Client','Month','Contrat'])
     ...: df1 = df.copy()
     ...: 
In [146]: %timeit original_app(df)
1 loops, best of 3: 645 ms per loop
In [147]: %timeit vectorized_app(df1)
100 loops, best of 3: 2.62 ms per loop

Answer 2

使用DataFrameGroupBy.size方法:

df.set_index(['Client', 'Month'], inplace=True)
df['Nbcontrats'] = df.groupby(level=(0,1)).size()
df.reset_index(inplace=True)

最重要的工作是将结果赋值回源DataFrame的列。