多向量的Numpy分组，获取分组索引

Question

我有几个numpy数组，我想构建一个分组方法，为这些数组分配组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)

- rinspy

Answer 1

在对数组 a 和 b 使用 np.stack 后，如果您在 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)

您可以用向量列表替换np.stack中的[a,b]。

编辑：更快的解决方案是在数组的sum上使用np.unique，乘以groupcols中所有先前数组的max加1的累积乘积(np.cumprod)。例如：

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.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)

但是组本身是相同的。

现在来检查时间:

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)

- Ben.T

---

谢谢，但我已经在原帖中添加了一些基准测试，看起来这种方法仍然比pandas实现慢7倍。 - rinspy

看我的更新 - 它们似乎差不多。Numpy 稍微慢一点，但没有实质性的区别。有趣的是，当长度为 3000 时，Numpy 实际上快了约 4 倍。 - rinspy

@rinspy 确实很有趣。我会看一下这个的。顺便说一句，如果你想用 pandas 来改进你的方法，你可以使用 df = pd.DataFrame({i:ar for i, ar in enumerate(groupcols)}).fillna(-1) 来创建数据框，它会稍微快一些 :) - Ben.T

@rinspy，如果你想看一下，我添加了一个更快的解决方案。 - Ben.T

这个numpy解决方案在“object”数据类型上不起作用... 看起来pandas的解决方案更加健壮。 - rinspy

Answer 2

numpy_indexed 包（免责声明：我是其作者）涵盖了以下类型的用例：

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

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

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

- Eelco Hoogendoorn

---

谢谢 - 我现在无法尝试它，但将其与pandas实现进行基准测试会很有趣 - 请参阅我的最新编辑。 - rinspy

它在性能特征方面类似于numpy实现。该实现是纯numpy的；我制作这个包的原因是为了用一个干净的接口封装这些操作。 - Eelco Hoogendoorn

Answer 3

我为另一个问题写了一个名为 group_by 的函数（这里）。这个函数非常灵活，可以解决你所要求的问题。

from itertools import count

# First test case:
a = np.array([1, 1, 1, 2, 2, 3])
b = np.array([1, 2, 2, 2, 3, 3])
data = np.stack([a, b], axis=-1)
ids = count()
print(group_by(data, lambda _: next(ids), transform=True)) # [0 1 1 2 3 4]

# Second test case:
a = np.array([1, 2, 1, 1, 1, 2, 3, 1])
b = np.array([1, 2, 2, 2, 2, 3, 3, 2])
data = np.stack([a, b], axis=-1)
ids = count()
second = group_by(data, lambda _: next(ids), transform=True)
print(second, second[[2, 3, 4, 7]]) # [0 2 1 1 1 3 4 1] [1 1 1 1]

我想指出，被接受的答案并没有回答你列举分组的问题。它计算的是逆索引。但它确实强调了np.unique的参数axis，这是最关键的一点。

- Carlos Pinzón