Python numpy多列分组

Question

Python numpy多列分组

3

有没有办法在numpy中通过多列进行分组聚合？我正在尝试使用这个模块：https://github.com/ml31415/numpy-groupies 目标是比pandas更快地进行分组聚合。例如：

group_idx = np.array([
np.array([4, 3, 3, 4, 4, 1, 1, 1, 7, 8, 7, 4, 3, 3, 1, 1]),
np.array([4, 3, 2, 4, 7, 1, 4, 1, 7, 8, 7, 2, 3, 1, 14 1]),
np.array([1, 2, 3, 4, 5, 1, 1, 2, 3, 4, 5, 4, 2, 3, 1, 1])
]
a = np.array([1, 2, 1, 2, 1, 2, 1, 2, 3, 4, 5, 4, 2, 3, 1, 1])

result = aggregate(group_idx, a, func='sum')

应该像这样：pandas df.groupby(['列1'，'列2'，'列3']).sum().reset_index()

- Christian Frei

group_idx中的数字都是正数吗？ - Divakar

是的，group_idx 中只有正值。 - Christian Frei

你是否同意添加对 numba 的依赖？ - senderle

当然，我认为这是可能的。 - Christian Frei

你能展示给我们精确的输出格式吗？你是想要带有索引和总和的二维数组输出吗？ - Divakar

是的，2D。在pandas中：df.groupby（['column1'，'column2'，'column3']）。sum（）。reset_index（） - Christian Frei

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Divakar · Accepted Answer

鉴于group_idx具有正值，我们可以使用基于降维的方法。我们假定前三列为分组列，最后一列（第四列）是要求和的数据列。

方法 #1

我们将坚持使用NumPy工具，并将pandas.factorize引入混合使用。

group_idx = df.iloc[:,:3].values
a = df.iloc[:,-1].values

s = group_idx.max(0)+1
lidx = np.ravel_multi_index(group_idx.T,s)

sidx, unq_lidx = pd.factorize(lidx)
pp = np.empty(len(unq_lidx), dtype=int)
pp[sidx] = np.arange(len(sidx))
k1 = group_idx[pp]

a_sums = np.bincount(sidx,a)
out = np.hstack((k1, a_sums.astype(int)[:,None]))

方法 #2

引入Numba和排序-

import numba as nb

@nb.njit
def step_sum(a_s, step_mask, out, group_idx_s):
    N = len(a_s)
    count_iter = 0
    for j in nb.prange(3):
        out[count_iter,j] = group_idx_s[0,j] 
    out[count_iter,3] = a_s[0]
    for i in nb.prange(1,N):
        if step_mask[i-1]:
            out[count_iter,3] += a_s[i]
        else:
            count_iter += 1
            for j in nb.prange(3):
                out[count_iter,j] = group_idx_s[i,j] 
            out[count_iter,3] = a_s[i]
    return out

group_idx = df.iloc[:,:3].values
a = df.iloc[:,-1].values

s = group_idx.max(0)+1
lidx = np.ravel_multi_index(group_idx.T,s)

sidx = lidx.argsort()
lsidx = lidx[sidx]
group_idx_s = group_idx[sidx]
a_s = a[sidx]

step_mask = lsidx[:-1] == lsidx[1:]    
N = len(lsidx)-step_mask.sum()
out = np.zeros((N, 4), dtype=int)
out = step_sum(a_s, step_mask, out, group_idx_s)

比较检查

为进行比较检查，我们可以使用类似以下的代码：

# get pandas o/p and lexsort
p = df.groupby(['agg_a','agg_b','agg_c'])['to_sum'].sum().reset_index().values
p = p[np.lexsort(p[:,:3].T)]

# Output from our approaches here, say `out`. Let's lexsort
out = out[np.lexsort(out[:,:3].T)]

print(np.allclose(out, p))