numpy 中的高效 bin 分配

Question

我有一个非常大的一维 python 数组 x，其中有些重复的数字以及一些相同大小的数据 d。

x = np.array([48531, 62312, 23345, 62312, 1567, ..., 23345, 23345])
d = np.array([0    , 1    , 2    , 3    , 4   , ..., 99998, 99999])

在我的上下文中 "very large" 指的是 10k...100k 条目。其中一些是重复的，因此唯一条目的数量约为 5k...15k。

我想将它们分组到垃圾箱中。这应该通过创建两个对象来完成。一个是矩阵缓冲区，b 的数据项取自 d。另一个对象是每个缓冲区列引用的唯一 x 值的向量 v。这是示例：

v =  [48531, 62312, 23345, 1567, ...]
b = [[0    , 1    , 2    , 4   , ...]
     [X    , 3    , ....., ...., ...]
     [ ...., ....., ....., ...., ...]
     [X    , X    , 99998, X   , ...]
     [X    , X    , 99999, X   , ...] ]

由于x中每个唯一数字的出现次数不同，缓冲区b中的某些值是无效的（由大写X表示，即"don't care"）。

在numpy中推导v非常容易：

v, n = np.unique(x, return_counts=True)  # yay, just 5ms

我们甚至得到 n，这是 b 中每列中有效条目的数量。而且，(np.max(n), v.shape[0])return是需要分配的矩阵b的shape。

但是如何高效地生成b呢？ for 循环可能会有所帮助

b = np.zeros((np.max(n), v.shape[0]))
for i in range(v.shape[0]):
    idx = np.flatnonzero(x == v[i])
    b[0:n[i], i] = d[idx]

此循环遍历 b 的所有列并通过识别 x == v.

的所有位置来提取索引 idx

但是我不喜欢这个解决方案，因为 for 循环相当慢（比 unique 命令长大约 50 倍）。我宁愿将操作向量化。

因此，一种矢量化方法是创建一个索引矩阵，其中 x == v 然后运行沿列在其上执行 nonzero() 命令。但是，此矩阵需要 150k x 15k 范围内的内存，因此在 32 位系统上大约需要 8GB。

对我来说，np.unique-操作甚至可以有效地 return 倒排索引，这样 x = v[inv_indices] 听起来很愚蠢，但是没有办法得到 v-to -x 为 v 中的每个 bin 分配列表。当函数扫描 x 时，这应该几乎是免费的。在实现方面，唯一的挑战是生成的索引矩阵的大小未知。

假设 np.unique-命令是用于分箱的方法来表述这个问题的另一种方式：

给定三个数组 x, v, inv_indices，其中 v 是 x 和 x = v[inv_indices] 中的唯一元素，是否有生成索引向量的有效方法 v_to_x[i] 这样 all(v[i] == x[v_to_x[i]]) 对于所有 bin i?

我不应该花比 np.unique 命令本身更多的时间。我很乐意为每个垃圾箱中的物品数量提供上限（例如 50）。

Answer 1

根据@user202729的建议我写了这段代码

x_sorted_args = np.argsort(x)
x_sorted = x[x_sorted_args]

i = 0
v = -np.ones(T)
b = np.zeros((K, T))

for k,g in groupby(enumerate(x_sorted), lambda tup: tup[1]):
    groups = np.array(list(g))[:,0]
    size = groups.shape[0]

    v[i] = k
    b[0:size, i] = d[x_sorted_args[groups]]
    i += 1

in 运行大约需要 100 毫秒，这会带来相当大的加速 w.r.t。上面发布的原始代码。

首先枚举出x中的值，加上相应的索引信息。然后枚举按实际 x 值分组，该值实际上是 enumerate().

生成的元组的第二个值

for 循环遍历所有组，将元组 g 的迭代器转换为大小为 (size x 2) 的 groups 矩阵，然后丢弃第二列，即 x 值仅保留索引。这导致 groups 只是一个一维数组。

groupby() 仅适用于排序数组。

干得好。我只是想知道我们是否可以做得更好？似乎仍然有很多不合理的数据复制发生。创建一个元组列表，然后将其转换为 2D 矩阵只是为了丢弃其中的一半仍然感觉有点次优。

Answer 2

我通过重新表述问题得到了我正在寻找的答案，请参阅此处：

by "cumulative counting" inv_indices 由 np.unique() 返回，我们收到稀疏矩阵的数组索引，因此

c = cumcount(inv_indices)
b[inv_indices, c] = d

上面链接的线程中提出的累积计数非常有效。运行低于 20 毫秒的时间非常现实。

numpy 中的高效 bin 分配

Efficient bin-assignment in numpy

python

indexing

numpy

binning

set-operations