更快的嵌套for循环,也许使用pd.groupby

时间:2019-07-03 07:54:40

标签: python for-loop nested pandas-groupby

我有一个嵌套循环。但是我的数据集非常大,所以我需要一种更快的方法。我相信可以通过某种巧妙的方式对数据进行分组或映射来完成。 我有一个X名称列表。对于每个名称,我都有包含颜色,品牌和值的项目。 我需要为每个名称的颜色和品牌的每种组合求和。然后,首先在每种颜色内,然后在每个名称内,对这些总和进行一些操作,同时跟踪计算总和的组合。 由于我的数据集很大,因此嵌套的for循环变得非常慢。 有没有更快的方法?

import random
import pandas
list1 = ['Name 0']
list2 = ['Color 0']
list3 = ['Brand 0']
list4 = [random.randint(10,1000)]
nName = 25
nColor = 5
nBrand = 3
nSim = 1000
for i in range(1,nSim):
    list1.extend(['Name {}'.format(random.randint(0,nName))])
    list2.extend(['Color {}'.format(random.randint(0,nColor))])
    list3.extend(['Brand {}'.format(random.randint(0, nBrand))])
    list4.append(random.randint(10,1000))
d = list(zip(list1, list2, list3, list4))
df1 = pd.DataFrame(d, columns=['Name','Color','Brand','Value'])

end_values = np.zeros(nName)
for iName in range(0,nName):
    y = 0
    for iColor in range(0,nColor):
        x = np.zeros(nBrand)
        for iBrand in range(0,nBrand):
            x[iBrand] = np.nansum(df1[((df1['Name'] == 'Name {}'.format(iName)) & (df1['Color'] == 'Color {}'.format(iColor)) & (df1['Brand'] == 'Brand {}'.format(iBrand)))]['Value'])
        y = y + x[0] + 1.5 * x[1] + 3 * x[0] * x[2]
    end_values[iName] = y

1 个答案:

答案 0 :(得分:0)

目前还不清楚您的问题是什么。因为问题中提供的代码中没有嵌套的for循环。您是指NameColorBrand值的所有组合上的嵌套循环吗?如果是这样,则无需编写任何循环。只需使用pandas意味着:

df2 = df1.groupby(['Name', 'Color', 'Brand']).agg({'Value': 'sum'})

是的,只需一行即可全部计算出来。就像SQL中的group by。例如,num22df2.loc[('Name 0', 'Color 2', 'Brand 2'), 'Value']

如果计算更加复杂

df3 = df1.groupby(['Name', 'Color', 'Brand']).agg({'Value': ['count', 'sum']})

num22df3.loc[('Name 0', 'Color 2', 'Brand 2'), ('Value', 'sum')]

.agg(...)的自变量是字典,您可以在其中指定要处理的列以及要获取的统计信息。

尝试这种方法,告诉我是否误解了您的问题。

编辑

考虑以下代码:

# 1. original code with nested loops
# includes slight correction (+1 in range(...) and zeros(...) functions)
def f_1(df1):
    end_values = np.zeros(nName + 1)
    for iName in range(nName + 1):
        y = 0
        for iColor in range(nColor + 1):
            x = np.zeros(nBrand + 1)
            for iBrand in range(nBrand + 1):
                x[iBrand] = np.nansum(df1[((df1['Name'] == 'Name {}'.format(iName)) & (df1['Color'] == 'Color {}'.format(iColor)) & (df1['Brand'] == 'Brand {}'.format(iBrand)))]['Value'])
            y = y + x[0] + 1.5 * x[1] + 3 * x[0] * x[2]
        end_values[iName] = y
    return end_values

# 3. new code - supportive function to calculate result for brands
def f_2_suppport(srs):
    cnt = {k: v for k, v in zip(srs.index.get_level_values(2), srs)}
    a = cnt.get('Brand 0', 0)
    b = cnt.get('Brand 1', 0)
    c = cnt.get('Brand 2', 0)
    return a + 1.5 * b + 3 * a * c

# 2. new code - main function with pandas groupby/agg
def f_2(df1):

    return (
        df1
        .groupby(['Name', 'Color', 'Brand'])
        .agg({'Value': 'sum'})
        .groupby(['Name', 'Color'])
        .agg({'Value': lambda x: f_2_suppport(x)})
        .groupby('Name')
        .agg({'Value': 'sum'}))


if __name__ == '__main__':

    print(f_1(df1))
    print(f_2(df1))

输出:

[28735057.  26111382.  28304153.5 15495841.  15915862.5 13884354.5
 21172873.  27342470.5 29403143.  38212958.5 37504605. ]
              Value
Name               
Name 0   28735057.0
Name 1   26111382.0
Name 10  37504605.0
Name 2   28304153.5
Name 3   15495841.0
Name 4   15915862.5
Name 5   13884354.5
Name 6   21172873.0
Name 7   27342470.5
Name 8   29403143.0
Name 9   38212958.5

我的机器上的运行时间:

%timeit f_1(df1)
1.07 s ± 736 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2(df1)
35.2 ms ± 163 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

因此groupby + agg使计算速度提高了约30倍。