我正在尝试向量化或以其他方式更快(可能使用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
数组的变换输入。
任何想法/指针?谢谢!
您可以使用布尔掩码来索引Sx
和prepped_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]
条件的元素。