返回数组中的索引1

Question

比方说，我们有一个类似[[1,0,0],[0,1,0],[0,0,1]]的numpy数组，我们希望将该数组更改为一个，其中ith项是原始数组的ith子数组中的索引1。在上述情况下，我们希望获得[0,1,2]作为返回值。

在numpy中有一个库方法可以做到这一点吗？

Answer 1

对于二进制（只有1和零）数组，可以使用numpy.argmax：

import numpy as np

arr = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
print(arr.argmax(axis=1))

输出

[0 1 2]

如果还有其他值，并且您想查找1，则可以使用：

arr = np.array([[1, 2, 3], [4, 1, 5], [6, 7, 1]])
print((arr == 1).argmax(axis=1))

输出

[0 1 2]

请注意，这也适用于二进制情况。