Python Numpy通过多个向量分组，获取组索引

我有几个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都应该是相同的

编辑2：添加一些基准

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)

---

在数组

a

和

b

上使用

np.stack

后，如果在中将参数

return\u 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)

您可以将

np.stack

中的

[a，b]

替换为所有向量的列表

编辑：一个更快的解决方案是对数组的

sum

使用

np.unique

乘以

max

的累积乘积（

np.cumprod

）加上

groupcols

中所有先前数组的1。例如：

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)

注意：与每个组关联的数字可能不相同（这里我将

a

的第一个元素更改为3）

但群体本身是一样的

现在检查时间：

a = np.random.randint(1, 100, 30000)
b = np.random.randint(1, 100, 30000)
c = np.random.randint(1, 100, 30000)
groupcols = [a,b,c]

%timeit group_pd(groupcols)
#13.7 ms ± 1.22 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit group_np2(groupcols)
#34.2 ms ± 6.88 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit group_np_sum(groupcols)
#3.63 ms ± 562 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

该软件包（dsiclaimer:I am it authos）涵盖了以下类型的用例：

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

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

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

---

谢谢，但我已经在原来的帖子中添加了一些基准测试，而且这个方法似乎仍然比pandas实现慢7倍。@rinspy只是一个问题，当您使用带有3个数组a、b和c的函数

group_pd

时，您是否更改了

groupby

列中的

groupby（[0，1，2]）

而不是

groupby（[0，1]）

或者您只想要前两个数组上的组？@rinspy确实很有趣。我来看看这个。顺便说一句，如果你想用

pandas

改进你的方法，你可以用

df=pd.dataframe（{i:ar for i，ar in enumerate（groupcols）}）创建dataframe。fillna（-1）

会稍微快一点：）@rinspy我添加了一个更快的解决方案，如果你想看看这个numpy解决方案对“object”数据类型不起作用。。。pandas解决方案似乎更健壮。谢谢-我目前无法尝试，但将其与pandas实现进行基准测试会很有趣-请参阅我的最新编辑。它在性能特征方面与numpy实现类似。实现是纯粹的numpy；我制作这个包的原因是将这些操作包装在一个干净的界面中。

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