python - 如何numpy-ify二维条件查找？

问题描述

我正在尝试矢量化或以其他方式更快（可能使用 numpy）查找/匹配 for 循环。我研究了np.vectorize、numpy 索引和np.where，但找不到适合我需要的正确实现/组合。

有问题的代码：

Sx = np.zeros((Np+1, 2*N+1))
rows, cols = prepped_array.shape[0], prepped_array.shape[1]

for ind1 in range(rows):
    for ind2 in range(cols):
        if prepped_array[ind1][ind2][0] != -1:
            Sx[ind1, ind2] = M[prepped_array[ind1][ind2][0], prepped_array[ind1][ind2][1]]

prepped_array是一个查找表（初始化为 all [-1, -1]），其中的值已被替换，它们应该在中更改Sx。

M是我们想要映射到Sx数组中的转换输入。

任何想法/指针？谢谢！

标签： pythonnumpy

您可以使用布尔掩码对进行索引Sx，prepped_array然后使用从派生的两个索引数组prepped_array对数组进行索引M。代码可能比上一句说得更清楚：

mask = prepped_array[:, :, 0] != -1
Sx[mask] = M[tuple(prepped_array[mask].T)]

让我们看看所涉及的步骤：

mask = prepped_array[:, :, 0] != -1创建一个二维布尔数组，指示满足条件的位置。
prepped_array[mask]创建一个 2D 数组，其中前 3 维的条目现在沿 2 维出现；第一个维度对应于中的每个True实例mask。
tuple(prepped_array[mask].T)创建两个可用于进一步索引其他数组的一维数组：第一个数组表示行索引，第二个数组表示列索引。
因此，使用前两个 1D 索引数组将Sx[mask] = M[tuple(prepped_array[mask].T)]包含的索引映射prepped_array到数组。M
Sx[mask]Sx最后引用那些满足条件的元素prepped_array[:, :, 0]。

python - 如何numpy-ify二维条件查找？

问题描述

解决方案

推荐阅读