Question

我有几个numpy数组；我想建立一个groupby方法，该方法具有这些数组的组ID。然后，它将允许我在组ID上索引这些数组，以对组执行操作。

例如：

import numpy as np
import pandas as pd
a = np.array([1,1,1,2,2,3])
b = np.array([1,2,2,2,3,3])

def group_np(groupcols):
    groupby = np.array([''.join([str(b) for b in bs]) for bs in zip(*[c for c in groupcols])])
    _, groupby = np.unique(groupby, return_invesrse=True)
   return groupby

def group_pd(groupcols):
    df = pd.DataFrame(groupcols[0])
    for i in range(1, len(groupcols)):
        df[i] = groupcols[i]
    for i in range(len(groupcols)):
        df[i] = df[i].fillna(-1)
    return df.groupby(list(range(len(groupcols)))).grouper.group_info[0]

输出：

group_np([a,b]) -> [0, 1, 1, 2, 3, 4]
group_pd([a,b]) -> [0, 1, 1, 2, 3, 4]

是否有一种更有效的实现方式，理想情况下是纯粹的numpy？目前，瓶颈似乎正在构建一个向量，该向量将为每个组具有唯一的值-目前，我正在通过将每个向量的值串联为字符串来实现这一目的。

我希望它适用于任意数量的输入向量，其中可以包含数百万个元素。

编辑：这是另一个测试用例：

a = np.array([1,2,1,1,1,2,3,1])
b = np.array([1,2,2,2,2,3,3,2])

这里，组元素2、3、4、7应该都相同。

Edit2：添加一些基准。

a = np.random.randint(1, 1000, 30000000)
b = np.random.randint(1, 1000, 30000000)
c = np.random.randint(1, 1000, 30000000)

def group_np2(groupcols):
    _, groupby = np.unique(np.stack(groupcols), return_inverse=True, axis=1)
    return groupby

%timeit group_np2([a,b,c])
# 25.1 s +/- 1.06 s per loop (mean +/- std. dev. of 7 runs, 1 loop each)
%timeit group_pd([a,b,c])
# 21.7 s +/- 646 ms per loop (mean +/- std. dev. of 7 runs, 1 loop each)

Answer 1

在数组np.stack和a上使用b之后，如果在np.unique中将参数return_inverse设置为True，则它是您正在寻找的输出：

a = np.array([1,2,1,1,1,2,3,1])
b = np.array([1,2,2,2,2,3,3,2])
_, inv = np.unique(np.stack([a,b]), axis=1, return_inverse=True)
print (inv)

array([0, 2, 1, 1, 1, 3, 4, 1], dtype=int64)

，您可以将[a,b]中的np.stack替换为所有向量的列表。

编辑：一种更快的解决方案是在数组np.unique上使用sum乘以{{1}的累积乘积（np.cumprod） }加上max中所有以前的数组中的1个。例如：

groupcols

要检查：

def group_np_sum(groupcols):
    groupcols_max = np.cumprod([ar.max()+1 for ar in groupcols[:-1]])
    return np.unique( sum([groupcols[0]] +
                          [ ar*m for ar, m in zip(groupcols[1:],groupcols_max)]), 
                      return_inverse=True)[1]

注意：与每个组关联的数字可能不相同（这里我将a = np.array([1,2,1,1,1,2,3,1]) b = np.array([1,2,2,2,2,3,3,2]) print (group_np_sum([a,b])) array([0, 2, 1, 1, 1, 3, 4, 1], dtype=int64)的第一个元素改为3）

但是组本身是相同的。

现在检查时间：

a = np.array([3,2,1,1,1,2,3,1])
b = np.array([1,2,2,2,2,3,3,2])
print(group_np2([a,b]))
print (group_np_sum([a,b]))
array([3, 1, 0, 0, 0, 2, 4, 0], dtype=int64)
array([0, 2, 1, 1, 1, 3, 4, 1], dtype=int64)

Answer 2

numpy_indexed程序包（dsiclaimer：我是它的授权人）涵盖了以下类型的用例：

import numpy_indexed as npi
npi.group_by((a, b))

像这样传递索引数组的元组可以避免创建副本。但是如果您不介意制作副本，也可以使用堆栈：

npi.group_by(np.stack(a, b))

多个向量对多个块进行求和，得到组索引

2 个答案: