加速计算附近的群组？

Question

我有一个包含组ID、两个距离度量（经度/纬度类型的度量）和一个值的数据框。对于给定的距离集，我想找到附近其他组的数量，以及这些附近其他组的平均值。

我编写了以下代码，但效率太低，对于非常大的数据集，它根本无法在合理的时间内完成。计算附近零售商很快，但计算附近零售商的平均值非常慢。有没有更好的方法使其更加高效？

distances = [1,2]

df = pd.DataFrame(np.random.randint(0,100,size=(100, 4)),
                  columns=['Group','Dist1','Dist2','Value'])

# get one row per group, with the two distances for each row
df_groups = df.groupby('Group')[['Dist1','Dist2']].mean()

# create KDTree for quick searching
tree = cKDTree(df_groups[['Dist1','Dist2']])

# find points within a given radius
for i in distances:
    closeby = tree.query_ball_tree(tree, r=i)

    # put into density column
    df_groups['groups_within_' + str(i) + 'miles'] = [len(x) for x in closeby]

    # get average values of nearby groups
    for idx, val in enumerate(df_groups.index):
        val_idx = df_groups.iloc[closeby[idx]].index.values
        mean = df.loc[df['Group'].isin(val_idx), 'Value'].mean()
        df_groups.loc[val, str(i) + '_mean_values'] = mean

    # merge back to dataframe
    df = pd.merge(df, df_groups[['groups_within_' + str(i) + 'miles', 
                                 str(i) + '_mean_values']], 
                  left_on='Group', 
                  right_index=True)

- user1566200

Answer 1

很明显，问题出在对主数据框进行索引时，使用了isin方法。随着数据框长度的增加，搜索范围也会变得更大。我建议您将同样的搜索操作应用在较小的df_groups数据框上，并计算更新后的平均值。

df = pd.DataFrame(np.random.randint(0,100,size=(100000, 4)),
                  columns=['Group','Dist1','Dist2','Value'])
distances = [1,2]
# get means of all values and count, the totals for each sample
df_groups = df.groupby('Group')[['Dist1','Dist2','Value']].agg({'Dist1':'mean','Dist2':'mean',
                                                                  'Value':['mean','count']})
# remove multicolumn index
df_groups.columns = [' '.join(col).strip() for col in df_groups.columns.values]
 #Rename columns 
df_groups.rename(columns={'Dist1 mean':'Dist1','Dist2 mean':'Dist2','Value mean':'Value','Value count':
                          'Count'},inplace=True)


# create KDTree for quick searching
tree = cKDTree(df_groups[['Dist1','Dist2']])

for i in distances:
    closeby = tree.query_ball_tree(tree, r=i)
    # put into density column
    df_groups['groups_within_' + str(i) + 'miles'] = [len(x) for x in closeby]
    #create column to look for subsets
    df_groups['subs'] = [df_groups.index.values[idx] for idx in closeby]
    #set this column to prep updated mean calculation
    df_groups['ComMean'] = df_groups['Value'] * df_groups['Count']

    #perform updated mean
    df_groups[str(i) + '_mean_values'] = [(df_groups.loc[df_groups.index.isin(row), 'ComMean'].sum() /
                                          df_groups.loc[df_groups.index.isin(row), 'Count'].sum()) for row in df_groups['subs']]
    df = pd.merge(df, df_groups[['groups_within_' + str(i) + 'miles',
                                 str(i) + '_mean_values']],
                  left_on='Group',
                  right_index=True)

更新后的平均值公式为(m1*n1 + m2*n2)/(n1+n2)，其中m1和m2是两个样本的平均值，n1和n2是它们的大小。

old setup 

100000 rows
%timeit old(df)
1 loop, best of 3: 694 ms per loop

1000000 rows
%timeit old(df)
1 loop, best of 3: 6.08 s per loop

10000000 rows
%timeit old(df)
1 loop, best of 3: 6min 13s per loop

新的设置

100000 rows
%timeit new(df)
10 loops, best of 3: 136 ms per loop

1000000 rows
%timeit new(df)
1 loop, best of 3: 525 ms per loop

10000000 rows
%timeit new(df)
1 loop, best of 3: 4.53 s per loop

- DJK