使用索引矩阵检索值的 Pythonic 方法?

我有一个 h-by-w-by-3 numpy 数组A和另一个 h-by-w-by-2 numpy 数组B,它告诉我A在每个 (h, w) 位置从哪里获取值。用C(h-by-w-by-3)表示结果矩阵。


天真,我可以


C = np.zeros(A.shape)

for i in range(B.shape[0]):

    for j in range(B.shape[1]):

        C[i, j, :] = A[B[i, j, 0], B[i, j, 1], :]

但我相信有更高效、更 Pythonic 和更快的方式!


慕盖茨4494581
浏览 127回答 1
1回答

跃然一笑

确实有:>>> import numpy as np>>> >>> h, w = 4, 6>>> >>> A = np.random.randint(0, 100, (h, w, 3))>>> B = np.random.randint(0, h*w, (h, w, 2)) % [h, w]>>> >>> C = np.zeros(A.shape)>>> for i in range(B.shape[0]):...     for j in range(B.shape[1]):...         C[i, j, :] = A[B[i, j, 0], B[i, j, 1], :]... >>> >>> C2 = A[(*np.moveaxis(B, 2, 0),)]>>> np.all(C == C2)True
打开App,查看更多内容
随时随地看视频慕课网APP

相关分类

Python