如何对二维条件查找进行 numpy-ify？

Question

我正在尝试矢量化或以其他方式使查找/匹配循环更快（可能使用 numpy）。我已经研究了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数组中的转换输入。

任何想法/指示？谢谢！

Answer 1

可以使用一个布尔掩码索引Sx和prepped_array，然后使用两个索引数组，从派生prepped_array，索引到M阵列。代码可能比上一句更清楚：mask = prepped_array[:, :, 0] != -1Sx[mask] = M[tuple(prepped_array[mask].T)]我们来看看涉及的步骤：mask = prepped_array[:, :, 0] != -1 创建一个二维布尔数组，指示满足条件的位置。prepped_array[mask]创建一个二维数组，其中来自先前第三维的条目现在沿着第二维出现；第一个维度对应于 中的每个True实例mask。tuple(prepped_array[mask].T) 创建两个一维数组，可用于进一步索引其他数组：第一个数组表示行索引，第二个数组表示列索引。因此Sx[mask] = M[tuple(prepped_array[mask].T)]，使用前两个一维索引数组将 中包含的索引映射prepped_array到M数组。Sx[mask]最后引用满足Sx条件的那些元素prepped_array[:, :, 0]。