Python numpy groupby 多列

Question

有没有办法通过 numpy 中的多列聚合来创建一个组？ 我试图用这个模块做到这一点：https://github.com/ml31415/numpy-groupies 目标是获得比 pandas 更快的 groupby。 例如：

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(['column1','column2','column3']).sum().reset_index()

Answer 1

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

方法#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))