python - 返回数组中 1 的索引

Question

假设我们有一个像这样的 numpy 数组[[1,0,0],[0,1,0],[0,0,1]]，我们想将该数组更改为第 i 个条目是原始数组的第 i 个子数组中 1 的索引的数组。在上述情况下，我们希望得到[0,1,2]作为返回值。

numpy 中是否有库方法可以做到这一点？

标签： pythonnumpy

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]

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