在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]])

Answer 2

列表理解怎么样？

ind = [[i,j] for i in range(3) for j in range(3) if A[i][j]>50]