以矢量化方式查找不同元素的索引

Question

我有一个ints的列表a，介于 0 到 3000 之间len(a) = 3000。我有一个for循环遍历这个列表，在更大的数组中搜索每个元素的索引。

import numpy as np

a = [i for i in range(3000)]

array = np.random.randint(0, 3000, size(12, 1000, 1000))

newlist = []

for i in range(0, len(a)):

    coord = np.where(array == list[i])

    newlist.append(coord)

如您所见，coord3D 矩阵中的值等于列表中的值的坐标 x、y、z 的 3 个数组。

for有没有办法在没有循环的情况下以矢量化方式做到这一点？

输出应该是一个元组列表，每个元素对应一个a：

# each coord looks like this:

print(coord)

(array[1, ..., 1000], array[2, ..., 1000], array[2, ..., 12])

# combined over all the iterations:

print(newlist)

[coord1, coord2, ..., coord3000]

Answer 1

尽管结果数组的大小都不同，但实际上有一个完全矢量化的解决方案。这个想法是这样的：对数组的所有元素及其坐标进行排序。argsort非常适合这种事情。在排序的数据中找到切点，这样您就知道在哪里拆分数组，例如使用diffand flatnonzero。split沿着您找到的索引的坐标数组。如果您缺少元素，您可能需要根据每次运行的第一个元素生成一个密钥。这是一个带您完成它的示例。假设您有一个dsize 的维数组n。您的坐标将是一个(d, n)数组：d = arr.ndim n = arr.size您可以直接生成坐标数组np.indices：coords = np.indices(arr.shape)现在ravel/reshape将数据和坐标分别放入一个(n,)和(d, n)数组中：arr = arr.ravel()  # Ravel guarantees C-order no matter the source of the data coords = coords.reshape(d, n)  # C-order by default as a result of `indices` too现在对数据进行排序：order = np.argsort(arr) arr = arr[order] coords = coords[:, order]查找数据更改值的位置。您需要新值的索引，因此我们可以制作一个比实际第一个元素小 1 的假第一个元素。change = np.diff(arr, prepend=arr[0] - 1)位置的索引给出了数组中的断点：locs = np.flatnonzero(change)您现在可以在这些位置拆分数据：result = np.split(coords, locs[1:], axis=1)您可以创建实际找到的值的键：key = arr[locs]如果您非常确信数组中存在所有值，那么您不需要密钥。相反，您可以计算locsas justnp.diff(arr)和resultas just np.split(coords, inds, axis=1)。中的每个元素result已经与where/使用的索引一致nonzero，但作为一个 numpy 数组。如果特别想要一个元组，您可以将其映射到一个元组：result = [tuple(inds) for inds in result]TL;博士将所有这些组合成一个函数：def find_locations(arr):    coords = np.indices(arr.shape).reshape(arr.ndim, arr.size)    arr = arr.ravel()    order = np.argsort(arr)    arr = arr[order]    coords = coords[:, order]    locs = np.flatnonzero(np.diff(arr, prepend=arr[0] - 1))    return arr[locs], np.split(coords, locs[1:], axis=1)您可以通过将最后一行替换为缺少元素的空数组返回索引数组列表    result = [np.empty(0, dtype=int)] * 3000   # Empty array, so OK to use same reference    for i, j in enumerate(arr[locs]):        result[j] = coords[i]    return result您可以选择过滤您想要的特定范围内的值（例如 0-2999）。

Answer 2

您可以在 numpy 中使用逻辑 OR 来一次传递所有这些相等条件，而不是一个一个地传递。import numpy as npconditions = Falsefor i in list:  conditions = np.logical_or(conditions,array3d == i)newlist = np.where(conditions)这允许 numpy 进行一次过滤，而不是分别为每个条件进行 n 次传递。另一种更紧凑的方法np.where(np.isin(array3d, list))