假设我有一个 N 维 numpy 数组 x 和一个 (N-1) 维索引数组 m(例如，m = x。 argmax(轴=-1))。我想构造 (N-1) 维数组 y 使得 y[i_1, ..., i_N-1] = x[i_1, ..., i_N-1 , m[i_1, ..., i_N-1]](对于上面的 argmax 示例，它等同于 y = x.max(axis=-1) )。
对于 N=3 我可以通过
y = x[np.arange(x.shape[0])[:, np.newaxis], np.arange(x.shape[1]), m]
问题是，我如何对任意 N 执行此操作？
最佳答案
你可以使用indices :
firstdims=np.indices(x.shape[:-1])
然后添加你的:
ind=tuple(firstdims)+(m,)
那么 x[ind] 就是你想要的。
In [228]: allclose(x.max(-1),x[ind])
Out[228]: True
关于python - numpy 中的索引(与 max/argmax 相关)，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/36315762/