在numpy矩阵中找到给定条件的指数

Question

所以我有一个numpy矩阵,例如:

[[1,2,3],
[0,59,2],
[54,4,2]]

现在我想找到索引,其中所有值都大于50(不是最大值).这应该给我[1,1],[2,0].

除了迭代,检查每个值并跟踪符合条件的值的索引然后返回这些对 - 你如何有效地做到这一点？

Answer 1

你想使用numpy.where或numpy.argwhere:

import numpy as np
A = np.array([[99, 2, 3],
              [0, 59, 2],
              [54, 4, 2]])
np.where(A > 50)
# (array([0, 1, 2]), array([0, 1, 0]))
np.argwhere(A > 50)
# array([[0, 0],
#        [1, 1],
#        [2, 0]])