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

Question

我有一个 int 的列表，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)

如您所见，coord 将是坐标 x、y、z 的 3 个数组，对应于 3D 矩阵中的值等于列表中的值。

有没有办法在没有 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

您可以在 numpy 中使用逻辑或来一次传递所有这些相等条件，而不是一个一个地传递。

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))

Answer 2

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

根据坐标对数组的所有元素进行排序。 argsort 非常适合这类事情。
在排序的数据中找到切点，这样你就知道在哪里拆分数组，例如diff and flatnonzero.
split 沿着您找到的索引的坐标数组。如果缺少元素，则可能需要根据每个运行.

这里有一个示例可以引导您完成它。假设您有一个大小为 n 的 d 维数组。您的坐标将是一个 (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]

如果您非常确信所有值都存在于数组中，那么您就不需要密钥。相反，您可以将 locs 计算为 np.diff(arr)，将 result 计算为 np.split(coords, inds, axis=1).

result 中的每个元素已经与 where/nonzero 使用的索引一致，但作为一个 numpy 数组。如果特别想要一个元组，你可以将它映射到一个元组：

result = [tuple(inds) for inds in result]

TL;DR

将所有这些组合成一个函数：

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)

您可以 return 通过将最后一行替换为

来创建包含缺失元素的空数组的索引数组列表

    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）。

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

Finding the indices of distinct elements in a vectorized manner

python

loops

numpy

vectorization