python - 从 numpy.argmax 结果中提取 numpy 数组切片

Question

我有一个 (3,3) numpy 数组，想找出绝对值的元素的符号值为最大值:

X = [[-2.1,  2,  3],
     [ 1, -6.1,  5],
     [ 0,  1,  1]]

s = numpy.argmax(numpy.abs(X),axis=0) 

给我所需元素的索引，s = [0,1,1]。

我如何使用此数组提取元素 [ -2.1, -6.1, 5] 以找出它们的符号？

Answer 1

试试这个:

# You might need to do this to get X as an ndarray (for example if X is a list)
X = numpy.asarray(X)

# Then you can simply do
X[s, [0, 1, 2]]

# Or more generally
X_argmax = X[s, numpy.arange(X.shape[1])]