首页 > 解决方案 > 蟒蛇 | NumPy - 将 y 值与相同 x 值相加的算法性能不佳

问题描述

我正在尝试开发一种算法,将每个 y 值与相同的 x 值相加。以下方法适用于小型数据集,但是一旦行数超过数万行,使用 for 循环对所有唯一 x 值求和非常慢。是否有另一种方法可以做到这一点,不涉及一次迭代每个唯一的 x 值一个值?

import numpy as np

data = np.array([[1, 2], [1, 3], [1, 5], [2, 2], [2, 4], [3, 1], [3, 8], [3, 9]])
x= data[:,0]
x = np.unique(x)
y = []
for item in x:
            y.append(data[data[:,0] == item].sum())
    
plt.plot(x,y)
plt.show()

标签: pythonnumpy

解决方案


使用 Dataframe 提高性能(100K 行的速度提高了约 182 倍)

代码

def use_df(data):
    ' Method based upon Datarame manipulations '
    df = pd.DataFrame(data)

    # Duplicate column 0 since will group based upon it and want to include it in sum
    df['groups'] = df[0]

    # Pandas series with desired values
    return df.groupby(0, as_index=True).sum().sum(axis=1)

def use_op(data):
    ' Based upon original post '
    x= data[:,0]
    x = np.unique(x)
    y = []
    for item in x:
        y.append(data[data[:,0] == item].sum())
    return y

验证

检查两个函数是否产生相同的输出。

 data = np.array([[1, 2], [1, 3], [1, 5], [2, 2], [2, 4], [3, 1], [3, 8], [3, 9]])
 print(use_df(data).tolist()==using_op(data))
 # Output: True (i.e. both produce values of  [13, 10, 27])

时间

使用原始数据进行计时

use_op   use_df
110 us   5.45 ms  
 

使用模拟数据的时间

Number Rows              use_op    use_df
1000                     5.18 ms   2.38 ms
10,000                   5.93 ms   53.5 ms
100,000                 14.6 ms    2.67 secs

带有模拟数据的时序代码

def create_data(n):
    ' Create random numpy array dimension n x 2'
    k = n //10      # Value variation is 10% of the number of rows        
    data = np.random.randint(1, k+1, n)
    return np.reshape(data, (-1, 2))

for n in [1000, 10000, 100000]:
    print(n)
    data = create_data(n)
    %timeit use_df(data)
    %timeit use_op(data)

概括

  • 小数据集的原始方法更快
  • DataFrame 在 10K 行时快 9 倍,在 100K 行时快 182 倍。

推荐阅读