python - Numpy:向量化 np.argwhere

Question

我在 numpy 中有以下数据结构:

import numpy as np

a = np.random.rand(267, 173) # dense img matrix
b = np.random.rand(199) # array of probability samples

我的目标是获取每个条目 i在b , 找到 a 中所有值的 x,y 坐标/索引位置那是<= i ，然后随机选择该子集中的值之一:

from random import randint

for i in b:
  l = np.argwhere(a <= i) # list of img coordinates where pixel <= i
  sample = l[randint(0, len(l)-1)] # random selection from `l`

这“有效”，但我想矢量化采样操作(即将 for 循环替换为 apply_along_axis 或类似的循环)。有谁知道如何做到这一点？任何建议将不胜感激!

Answer 1

你不能完全向量化 np.argmax因为你每次都有一个随机的子集大小。不过，您可以做的是通过排序显着加快计算速度。对图像进行一次排序将创建一个分配，而在每一步对图像进行屏蔽将为提取的元素创建一个临时数组，用于屏蔽和。使用排序图像，您可以应用 np.searchsorted获取尺寸:

a_sorted = np.sort(a.ravel())
indices = np.searchsorted(a_sorted, b, side='right')

你仍然需要一个循环来进行采样，但你可以做类似的事情

samples = np.array([a_sorted[np.random.randint(i)] for i in indices])

获取 x-y 坐标而不是样本值对于这个系统来说有点复杂。您可以使用 np.unravel_index要获取索引，但首先您必须将 a_sorted 的引用框架转换为 a.ravel()。如果您使用 np.argsort 排序而不是 np.sort ，您可以获取原始数组中的索引。幸运的是，np.searchsorted 通过 sorter 参数支持这种情况:

a_ind = np.argsort(a, axis=None)
indices = np.searchsorted(a.ravel(), b, side='right', sorter=a_ind)
r, c = np.unravel_index(a_ind[[np.random.randint(i) for i in indices]], a.shape)

r和c与b大小相同，对应a中的行索引和列索引> 每个选择基于 b。索引转换取决于数组中的步幅，因此我们假设您使用的是 C 顺序，因为默认情况下 90% 的数组都会这样做。

复杂性

假设 b 的大小为 M，a 的大小为 N。

您当前的算法针对 b 的每个元素对 a 的每个元素进行线性搜索。在每次迭代中，它为匹配元素分配一个掩码(平均 N/2)，然后分配一个相同大小的缓冲区来保存掩码选择。这意味着时间复杂度在 O(M * N) 数量级，空间复杂度相同。

我的算法首先对a进行排序，这是O(N log N)。然后它搜索 M 个插入点，即 O(M log N)。最后，它选择 M 个样本。它分配的空间是图像的一个排序副本和两个大小为 M 的数组。因此，它的时间复杂度为 O((M + N) log N)，空间复杂度为 O(M + N)。