如何numpy-ify二维条件查找?

问题描述 投票:2回答:1

我正在尝试向量化或以其他方式更快(可能使用numpy)查找/匹配循环。我查看了np.vectorize,numpy indexing和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是一个查找表(初始化为所有[-1, -1]),其中值已被替换,应在Sx中更改它们。

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

任何想法/指针?谢谢!

python numpy
1个回答
0
投票

您可以使用布尔掩码来索引Sxprepped_array,然后使用从prepped_array派生的两个索引数组来索引M数组。代码可能比前一句更清晰:

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

让我们看一下所涉及的步骤:

  • mask = prepped_array[:, :, 0] != -1创建一个2D布尔数组,指示满足条件的位置。
  • prepped_array[mask]创建一个二维数组,其中前一个第三维的条目现在沿着第二维出现;第一个维度对应于True中的每个mask实例。
  • tuple(prepped_array[mask].T)创建了两个1D数组,可用于进一步索引到其他数组:第一个数组表示行索引,第二个数组表示列索引。
  • 所以Sx[mask] = M[tuple(prepped_array[mask].T)]使用前两个1D索引数组将prepped_array中包含的索引映射到数组M
  • Sx[mask]最终引用了Sx中那些符合prepped_array[:, :, 0]条件的元素。
© www.soinside.com 2019 - 2024. All rights reserved.