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

时间:2018-08-30 16:09:31

标签: python python-3.x pandas numpy

我的pandas / numpy生锈了,我编写的代码效率低下。

我正在Python3.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具有两列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 个答案:

答案 0 :(得分:4)

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)

答案 1 :(得分: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 为了后代,我附上@piRSquared关于numba为什么在这里有帮助的出色评论:

  

numba的优点是循环效率很高。虽然可以   了解NumPy的许多API,通常最好避免创建   循环内的NumPy对象。我的代码正在为创建一个NumPy数组   数据框中的每一行。然后在使用之前将它们串联   Bincount。 @jpp的numba代码创建的额外对象很少,   利用已经存在的很多东西。我之间的区别   NumPy解决方案和@jpp的numba解决方案约为4-5倍。两者都是   线性的,应该很快。

答案 2 :(得分:3)

我的解决方案

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