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

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)

如您所见，coord 将是 3 个坐标 x、y、z 的数组，用于 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

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

1. 对数组中的所有元素及其坐标进行排序。 argsort 非常适合这类事情。
2. 在排序后的数据中找到切点，这样您就知道在哪里拆分数组，例如diff 和 flatnonzero。
3. 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)

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

来返回索引数组列表，其中包含缺失元素的空数组

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