python - numpy 中的高效 bin 分配

Question

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

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

在我的上下文中，“非常大”是指 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 表示，即“不关心”)。

---

在 numpy 中推导 v 非常容易:

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

我们甚至得到 n，它是 b 中每一列中有效条目的数量。此外，(np.max(n), v.shape[0]) 返回需要分配的矩阵 b 的形状。

但是如何高效地生成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 操作甚至可以有效地返回倒排索引，这样 x = v[inv_indices] 听起来相当愚蠢，但是没有获取 v 中每个 bin 的 v 到 x 分配列表的方法。当函数扫描 x 时，这应该几乎是免费的。在实现方面，唯一的挑战是生成的索引矩阵的大小未知。

---

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

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

我不应该花比 np.unique-command 本身更多的时间。我很乐意为每个箱子中的元素数量提供上限(例如 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 矩阵只是为了扔掉它的一半仍然感觉有点次优。