Question

如何使用ndarray矩阵选择ndarray的元素？

这是我的意思的一个例子。

a = np.arange(9)
b = np.arange(5)
c = np.arange(12)
A, B, C = np.meshgrid(a, b, c, indexing='ij')

现在，对于a，c的每个值，我希望最小化b的{{1}}。获取指数：

A+C=B

显然，idx = np.abs(A+C-B).argmin(axis=1)的形状为idx - 它包含每个(9, 12) a和每个9 c的b索引。

现在，我想用＆＃34;优化的b＆＃34;选择矩阵。这就是

的内容

据说具有B[:, idx, :]形状 - 因为对于每个其他组合，它只有一个值(9, 1, 12) - 最小化一个。{1}}。现在，b改为为我提供了形状为B[:, idx, :]的所有潜在组合的网格。我也试过

(9, 9, 12, 12)

如何获得上述特定类型的矩阵？

Answer 1

您只需在np.newaxis/None处添加一个轴即可触发advanced-indexing -

B[np.arange(B.shape[0])[:,None], idx, np.arange(B.shape[2])]

这个想法主要是将idx的行映射到np.arange(B.shape[0])的第一个索引数组，因此我们需要在那里添加一个轴。为了映射idx的列，我们已经np.arange(B.shape[2])沿着列对齐。

替代np.newaxis

添加新轴的另一种方法是使用reshaping因此，我们可以将B[np.arange(B.shape[0])[:,None]替换为np.arange(B.shape[0]).reshape(-1,1)。

进一步优化

我们可以通过使用开放数组来替换由meshgrid创建的巨大数组来优化代码，就像这样 -

A0, B0, C0 = np.ix_(a,b,c)
idx = np.abs(A0+C0-B0).argmin(axis=1)

因此，得到最终输出，如此 -

B[np.arange(len(a))[:,None], idx, np.arange(len(c))]

只是为了让自己在这里节省内存的想法 -

In [47]: A.nbytes + B.nbytes + C.nbytes
Out[47]: 12960

而A0, B0, C0分别是输入数组a, b, c的视图，因此不会占用任何额外的内存，即完全免费 -

In [49]: np.shares_memory(a,A0)
Out[49]: True

为了完整起见，获得idx的直接方法是 -

np.abs(a[:,None,None]+c-b[:,None]).argmin(axis=1)