如何有效地迭代pandas DataFrame并在这些值上增加NumPy数组？

Question

我的熊猫/ numpy生锈了,我写的代码感觉效率低下.

我正在Python3.x初始化一个numpy零的数组,长度为1000.为了我的目的,这些只是整数:

import numpy as np
array_of_zeros =  np.zeros((1000, ), )

我还有以下DataFrame(比我的实际数据小得多)

import pandas as pd
dict1 = {'start' : [100, 200, 300], 'end':[400, 500, 600]}
df = pd.DataFrame(dict1)
print(df)
##
##    start     end
## 0    100     400
## 1    200     500
## 2    300     600

DataFrame有两列,start和end.这些值表示一系列值,即start总是小于的整数end.在上面,我们看到第一行有范围100-400,接下来是200-500,然后300-600.

我的目标是逐行遍历pandas DataFrame,并array_of_zeros根据这些索引位置递增numpy数组.因此,如果10to 的数据帧中有一行20,我想将索引10-20的零增加+1.

这是我想要的代码:

import numpy as np
array_of_zeros =  np.zeros((1000, ), )

import pandas as pd
dict1 = {'start' : [100, 200, 300], 'end':[400, 500, 600]}
df = pd.DataFrame(dict1)
print(df)

for idx, row in df.iterrows():
    for i in range(int(row.start), int(row.end)+1):
        array_of_zeros[i]+=1

它的工作原理!

print(array_of_zeros[15])
## output: 0.0
print(array_of_zeros[600])
## output: 1.0
print(array_of_zeros[400])
## output: 3.0
print(array_of_zeros[100])
## output: 1.0
print(array_of_zeros[200])
## output: 2.0

我的问题:这是非常笨拙的代码!我不应该使用那么多带有numpy数组的for循环!如果输入数据帧非常大,则此解决方案效率非常低

是否有更有效(即更多基于numpy)的方法来避免这种for循环？

for i in range(int(row.start), int(row.end)+1):
    array_of_zeros[i]+=1

也许有一个以熊猫为导向的解决方案？

Answer 1

您可以使用 NumPy 数组索引来避免内部循环，即res[np.arange(A[i][0], A[i][1]+1)] += 1，但这效率不高，因为它涉及创建新数组并使用高级索引。

\n\n

相反，您可以使用numba¹来优化您的算法，就像现在一样。下面的示例显示了通过将性能关键逻辑转移到 JIT 编译的代码来实现巨大的性能改进。

\n\n

from numba import jit\n\n@jit(nopython=True)\ndef jpp(A):\n    res = np.zeros(1000)\n    for i in range(A.shape[0]):\n        for j in range(A[i][0], A[i][1]+1):\n            res[j] += 1\n    return res\n

\n\n

一些基准测试结果：

\n\n

# Python 3.6.0, NumPy 1.11.3\n\n# check result the same\nassert (jpp(df[['start', 'end']].values) == original(df)).all()\nassert (pir(df) == original(df)).all()\nassert (pir2(df) == original(df)).all()\n\n# time results\ndf = pd.concat([df]*10000)\n\n%timeit jpp(df[['start', 'end']].values)  # 64.6 \xc2\xb5s per loop\n%timeit original(df)                      # 8.25 s per loop\n%timeit pir(df)                           # 208 ms per loop\n%timeit pir2(df)                          # 1.43 s per loop\n

\n\n

用于基准测试的代码：

\n\n

def original(df):\n    array_of_zeros = np.zeros(1000)\n    for idx, row in df.iterrows():\n        for i in range(int(row.start), int(row.end)+1):\n            array_of_zeros[i]+=1   \n    return array_of_zeros\n\ndef pir(df):\n    return np.bincount(np.concatenate([np.arange(a, b + 1) for a, b in \\\n                       zip(df.start, df.end)]), minlength=1000)\n\ndef pir2(df):\n    a = np.zeros((1000,), np.int64)\n    for b, c in zip(df.start, df.end):\n        np.add.at(a, np.arange(b, c + 1), 1)\n    return a\n

\n\n\n\n

¹对于后代，我添加了 @piRSquared 的精彩评论，说明为什么numba在这里有帮助：

\n\n

\n

numba的优点是循环非常高效。尽管它可以理解 NumPy 的大部分 API，但通常最好避免在循环中创建 NumPy 对象。我的代码正在为数据框中的每一行创建一个 NumPy 数组。然后在使用 bincount 之前将它们连接起来。@jpp 的numba代码创建了很少的额外对象，并利用了大部分已有的对象。我的\n NumPy 解决方案和@jpp 的numba解决方案之间的差异约为 4-5 倍。两者都是线性的并且应该相当快。

\n

\n