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

Question

我不太擅长使用 pandas/numpy，我写的代码感觉效率低下。

我正在 Python 3.x 中初始化一个长度为 1000 的 numpy 数组，并将其全部赋值为零。对于我的需求来说，这些只是整数：

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，并根据这些索引位置增加numpy数组array_of_zeros。因此，如果数据框中有一个10到20的行，我想将0增加+1，以便索引为10-20。

以下是实现我所需功能的代码：

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

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

- ShanZhengYang

Answer 1

numpy.bincount

np.bincount(np.concatenate(
    [np.arange(a, b + 1) for a, b in zip(df.start, df.end)]
), minlength=1000)

---

numpy.add.at

a = np.zeros((1000,), np.int64)
for b, c in zip(df.start, df.end):
  np.add.at(a, np.arange(b, c + 1), 1)

- piRSquared

Answer 2

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

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

from numba import jit

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

一些基准测试结果：

# Python 3.6.0, NumPy 1.11.3

# check result the same
assert (jpp(df[['start', 'end']].values) == original(df)).all()
assert (pir(df) == original(df)).all()
assert (pir2(df) == original(df)).all()

# time results
df = pd.concat([df]*10000)

%timeit jpp(df[['start', 'end']].values)  # 64.6 µs per loop
%timeit original(df)                      # 8.25 s per loop
%timeit pir(df)                           # 208 ms per loop
%timeit pir2(df)                          # 1.43 s per loop

用于基准测试的代码：

def original(df):
    array_of_zeros = np.zeros(1000)
    for idx, row in df.iterrows():
        for i in range(int(row.start), int(row.end)+1):
            array_of_zeros[i]+=1   
    return array_of_zeros

def pir(df):
    return np.bincount(np.concatenate([np.arange(a, b + 1) for a, b in \
                       zip(df.start, df.end)]), minlength=1000)

def pir2(df):
    a = np.zeros((1000,), np.int64)
    for b, c in zip(df.start, df.end):
        np.add.at(a, np.arange(b, c + 1), 1)
    return a

---

¹ 为了后人的参考，我在这里包含 @piRSquared 的优秀评论，解释了为什么 numba 对此有所帮助：

numba 可以非常高效地循环。虽然它可以理解 NumPy 的大部分 API，但通常最好避免在循环内创建 NumPy 对象。我的代码为每行数据框创建一个 NumPy 数组。然后在使用 bincount 之前将它们连接起来。@jpp 的 numba 代码几乎不会创建额外的对象，并且利用了已经存在的很多东西。我的 NumPy 解决方案和 @jpp 的 numba 解决方案之间的差异约为 4-5 倍。两者都是线性的，应该非常快。

- jpp

---

你知道我喜欢“numba”。另外，“from numba import njit”-> njit 是 jit（nopython=True）。 - piRSquared

@piRSquared，谢谢！我希望你不介意我从你的解决方案中添加时间。是的，我知道njit，但我认为对于那些刚接触numba的人来说，有时更明确一些。 - jpp

谢谢！jpp实现非常出色！我仍有一个问题：是否可能使用numba来优化@piRSquared的代码？ - ShanZhengYang

@ShanZhengYang numba的优势在于循环非常高效。虽然它可以理解大部分Numpy的api，但通常最好避免在循环中创建Numpy对象。我的代码为数据框中的每一行创建了一个Numpy数组，然后在使用bincount之前将它们连接起来。@jpp的numba代码几乎不会创建额外的对象，并且利用了已有的许多内容。我使用Numpy的解决方案和jpp的numba解决方案之间的差异约为4-5倍。两者都是线性的，应该非常快。 - piRSquared

为什么你使用了res[np.arange(A[i][0], A[i][1]+1)] += 1（分配一个临时数组，在循环中填充值，然后在另一个循环中使用它来索引数组res）？是否有误导性的教程需要进行更正？一个简单的嵌套循环由于上述明显的原因大约快2倍。尽管如此，即使对于这种方法，您的计时也相当慢。您是否在使用timeit之前调用了该函数一次？如果没有，您测量了第一次编译的恒定开销，大约为0.2秒，这也可以通过cache=True来最小化。 - max9111

Answer 3

我的解决方案

for x, y in zip(df.start, df.end):
    array_of_zeros[x:y+1]+=1

- BENY

---

array_of_zeros[x:y] += 1? - piRSquared

@piRSquared 我之前尝试过，似乎不起作用..让我再测试一下。 - BENY

你应该也不包括 y 吗？也就是说 x:y 排除了 y，难道不应该是 x:(y+1) 吗？ - Onyambu

@Onyambu 你好 double ?? 已经修复 :-) - BENY

使用循环是一种可行的方法。但它并不高效，需要更多的时间。请参阅此博客文章，了解有关迭代DataFrame及其效率和时间估计的不同方法的分析：https://towardsdatascience.com/apply-function-to-pandas-dataframe-rows-76df74165ee4 - Ahwar