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

Question

我有一个ints的列表a，介于 0 到 3000 之间len(a) = 3000。我有一个for循环遍历这个列表，在更大的数组中搜索每个元素的索引。import numpy as npa = [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

尽管结果数组的大小都不同，但实际上有一个完全矢量化的解决方案。这个想法是这样的：

1. 对数组的所有元素及其坐标进行排序。argsort非常适合这种事情。

2. 在排序的数据中找到切点，这样您就知道在哪里拆分数组，例如使用diffand flatnonzero。

3. 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 np

conditions = False

for i in list:

  conditions = np.logical_or(conditions,array3d == i)

newlist = np.where(conditions)

这允许 numpy 进行一次过滤，而不是分别为每个条件进行 n 次传递。

另一种更紧凑的方法

np.where(np.isin(array3d, list))