首页 文章

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

提问于
浏览
6

我的熊猫/ 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有两列, startend . 这些值表示一系列值,即 start 将始终是小于 end 的整数 . 在上面,我们看到第一行的范围是 100-400 ,接下来是 200-500 ,然后是 300-600 .

我的目标是逐行遍历pandas DataFrame,并根据这些索引位置递增numpy数组 array_of_zeros . 因此,如果 1020 的数据帧中有一行,我想将索引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

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

3 回答

  • 4

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

    相反,您可以使用 numba 1来优化您的算法,完全符合它的原样 . 以下示例通过将性能关键逻辑移至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
    

    1对于后人,我'm including @piRSquared'对 numba 在这里帮助的原因有很好的评论:

    numba的优势在于非常有效地循环 . 虽然它可以理解NumPy的大部分API,但通常最好避免在循环中创建NumPy对象 . 我的代码是为数据帧中的每一行创建一个NumPy数组 . 然后在使用bincount之前连接它们 . @jpp的numba代码创建了很少的额外对象,并利用了已有的大部分内容 . 我的NumPy解决方案和@jpp的numba解决方案之间的差异大约是4-5倍 . 两者都是线性的,应该很快 .

  • 3

    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)
    
  • 4

    我的解决方案

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

相关问题