如何矢量化（利用 pandas/numpy）而不是使用嵌套的 for 循环

Question

我希望有效地使用pandas（或numpy）而不是带有if 语句的嵌套for 循环来解决特定问题。这是一个玩具版本：

假设我有以下两个 DataFrames

import pandas as pd
import numpy as np

dict1 = {'vals': [100,200], 'in': [0,1], 'out' :[1,3]}
df1 = pd.DataFrame(data=dict1)

dict2 = {'vals': [500,800,300,200], 'in': [0.1,0.5,2,4], 'out' :[0.5,2,4,5]}
df2 = pd.DataFrame(data=dict2)

现在我希望遍历每个数据帧的每一行，并在满足特定条件时将 vals 相乘。这段代码适用于我想要的

ans = []

for i in range(len(df1)):
    for j in range(len(df2)):
        if (df1['in'][i] <= df2['out'][j] and df1['out'][i] >= df2['in'][j]):
            ans.append(df1['vals'][i]*df2['vals'][j])

np.sum(ans)

但是，显然这非常效率低下，实际上我的 DataFrame 可能有数百万个条目，这使得它无法使用。我也没有让我们使用pandas 或numpy 高效的向量实现。有谁知道如何有效地矢量化这个嵌套循环？

我觉得这段代码类似于矩阵乘法，所以使用outer 可以取得进展吗？这是if 条件，我发现很难融入其中，因为if 逻辑需要将df1 中的每个条目与df2 中的所有条目进行比较。

Answer 1

您也可以使用像 Numba 这样的编译器来完成这项工作。这也将优于矢量化解决方案，并且不需要临时数组。

示例

import numba as nb
import numpy as np
import pandas as pd
import time

@nb.njit(fastmath=True,parallel=True,error_model='numpy')
def your_function(df1_in,df1_out,df1_vals,df2_in,df2_out,df2_vals):
  sum=0.
  for i in nb.prange(len(df1_in)):
      for j in range(len(df2_in)):
          if (df1_in[i] <= df2_out[j] and df1_out[i] >= df2_in[j]):
              sum+=df1_vals[i]*df2_vals[j]
  return sum

测试

dict1 = {'vals': np.random.randint(1, 100, 1000),
         'in': np.random.randint(1, 10, 1000),
         'out': np.random.randint(1, 10, 1000)}
df1 = pd.DataFrame(data=dict1)
dict2 = {'vals': np.random.randint(1, 100, 1500),
         'in': 5*np.random.random(1500),
         'out': 5*np.random.random(1500)}
df2 = pd.DataFrame(data=dict2)

# First call has some compilation overhead
res=your_function(df1['in'].values, df1['out'].values, df1['vals'].values,
                  df2['in'].values, df2['out'].values, df2['vals'].values)

t1 = time.time()
for i in range(1000):
  res = your_function(df1['in'].values, df1['out'].values, df1['vals'].values,
                      df2['in'].values, df2['out'].values, df2['vals'].values)

print(time.time() - t1)

时间

vectorized solution @AGN Gazer: 9.15ms
parallelized Numba Version: 0.7ms

Answer 2

m1 = np.less_equal.outer(df1['in'], df2['out']) 
m2 = np.greater_equal.outer(df1['out'], df2['in'])
m = np.logical_and(m1, m2)
v12 = np.outer(df1['vals'], df2['vals'])
print(v12[m].sum())

或者，用这个长行替换前三行：

m = np.less_equal.outer(df1['in'], df2['out']) & np.greater_equal.outer(df1['out'], df2['in'])
s = np.outer(df1['vals'], df2['vals'])[m].sum()

对于非常大的问题，推荐dask。

计时测试：

这是使用 1000 和 1500 长数组时的时序比较：

In [166]: dict1 = {'vals': np.random.randint(1,100,1000), 'in': np.random.randint(1,10,1000), 'out': np.random.randint(1,10,1000)}
     ...: df1 = pd.DataFrame(data=dict1)
     ...: 
     ...: dict2 = {'vals': np.random.randint(1,100,1500), 'in': 5*np.random.random(1500), 'out': 5*np.random.random(1500)}
     ...: df2 = pd.DataFrame(data=dict2)

作者原创方法（Python循环）：

In [167]: def f(df1, df2):
     ...:     ans = []
     ...:     for i in range(len(df1)):
     ...:         for j in range(len(df2)):
     ...:             if (df1['in'][i] <= df2['out'][j] and df1['out'][i] >= df2['in'][j]):
     ...:                 ans.append(df1['vals'][i]*df2['vals'][j])
     ...:     return np.sum(ans)
     ...: 
     ...: 

In [168]: %timeit f(df1, df2)
47.3 s ± 1.02 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

@Ben.T 方法：

In [170]: %timeit df2['ans']= df2.apply(lambda row: df1['vals'][(df1['in'] <= row['out']) & (df1['out'] >= row['in'])].sum()*row['vals'],1); df2['a
     ...: ns'].sum()
2.22 s ± 40.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

这里提出的矢量化解决方案：

In [171]: def g(df1, df2):
     ...:     m = np.less_equal.outer(df1['in'], df2['out']) & np.greater_equal.outer(df1['out'], df2['in'])
     ...:     return np.outer(df1['vals'], df2['vals'])[m].sum()
     ...: 
     ...: 

In [172]: %timeit g(df1, df2)

7.81 ms ± 127 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answer 3

你的答案：

471 µs ± 35.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

方法1（慢3倍以上）：

df1.apply(lambda row: list((df2['vals'][(row['in'] <= df2['out']) & (row['out'] >= df2['in'])] * row['vals'])), axis=1).sum()

1.56 ms ± 7.56 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

方法2（慢2倍以上）：

ans = []
for name, row in df1.iterrows():
    _in = row['in']
    _out = row['out']
    _vals = row['vals']
    ans.append(df2['vals'].loc[(df2['in'] <= _out) & (df2['out'] >= _in)].values * _vals)

1.01 ms ± 8.21 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

方法3（快3倍以上）：

df1_vals = df1.values
ans = np.zeros(shape=(len(df1_vals), len(df2.values)))
for i in range(df1_vals.shape[0]):
    df2_vals = df2.values
    df2_vals[:, 2][~np.logical_and(df1_vals[i, 1] >= df2_vals[:, 0], df1_vals[i, 0] <= df2_vals[:, 1])] = 0
    ans[i, :] = df2_vals[:, 2] * df1_vals[i, 2]

144 µs ± 3.11 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

在方法 3 中，您可以通过执行以下操作查看解决方案：

ans[ans.nonzero()]

Out[]: array([ 50000.,  80000., 160000.,  60000.]

我想不出一种方法来删除底层循环:(但我在这个过程中学到了很多关于 numpy 的知识！（学习了）

Answer 4

一种方法是使用apply。在 df2 中创建一列，其中包含 df1 中的 val 总和，满足您的输入和输出标准，乘以 df2 行的 vals

df2['ans']= df2.apply(lambda row: df1['vals'][(df1['in'] <= row['out']) & 
                                              (df1['out'] >= row['in'])].sum()*row['vals'],1)

那么就对这一列求和

df2['ans'].sum()