python - 如何在满足多个条件的 numpy 数组中找到索引？

Question

我在 Python 中有一个这样的数组:

例子:

>>> scores = numpy.asarray([[8,5,6,2], [9,4,1,4], [2,5,3,8]])
>>> scores
array([[8, 5, 6, 2],
   [9, 4, 1, 4],
   [2, 5, 3, 8]])

我想在 scores 中找到所有 [row, col] 索引，其中的值为:

1) 行中的最小值

2) 大于阈值

3) 最多是行中下一个最大值的 .8 倍

我想尽可能高效地完成它，最好没有任何循环。我已经为此苦苦挣扎了一段时间，所以如果您能提供任何帮助，我将不胜感激!

Answer 1

它应该按照以下方式行事

In [1]: scores = np.array([[8,5,6,2], [9,4,1,4], [2,5,3,8]]); threshold = 1.1; scores
Out[1]: 
array([[8, 5, 6, 2],
       [9, 4, 1, 4],
       [2, 5, 3, 8]])

In [2]: part = np.partition(scores, 2, axis=1); part
Out[2]: 
array([[2, 5, 6, 8],
       [1, 4, 4, 9],
       [2, 3, 5, 8]])

In [3]: row_mask = (part[:,0] > threshold) & (part[:,0] <= 0.8 * part[:,1]); row_mask
Out[3]: array([ True, False,  True], dtype=bool)

In [4]: rows = row_mask.nonzero()[0]; rows
Out[4]: array([0, 2])

In [5]: cols = np.argmin(scores[row_mask], axis=1); cols
Out[5]: array([3, 0])

在那一刻，如果你正在寻找实际的坐标对，你可以zip它们:

In [6]: coords = zip(rows, cols); coords
Out[6]: [(0, 3), (2, 0)]

或者如果您打算查找这些元素，您可以按原样使用它们:

In [7]: scores[rows, cols]
Out[7]: array([2, 2])