python - 如果嵌套数组的最大值高于阈值，则获取嵌套数组的 Numpy 条件

Question

我有以下数组:

arr = numpy.array([[.5, .5], [.9, .1], [.8, .2]])

我想要获取 arr 的索引，其中包含一个数组，其最大值大于或等于 .9。因此，对于这种情况，结果将是 [1]，因为索引为 1 [.9, .1] 的数组是唯一一个最大值 >= 9.

我试过:

>>> condition = np.max(arr) >= .9
>>> arr[condition]
array([ 0.5,  0.5])

但是，如您所见，它会产生错误的答案。

Answer 1

我想你想要np.where这里。此函数返回满足特定条件的任何值的索引:

>>> np.where(arr >= 0.9)[0] # here we look at the whole 2D array
array([1])

(np.where(arr >= 0.9) 返回索引数组的元组，数组的每个轴一个。您的预期输出意味着您只需要行索引(轴 0 ).)

如果想先取每一行的最大值，可以使用arr.max(axis=1):

>>> np.where(arr.max(axis=1) >= 0.9)[0] # here we look at the 1D array of row maximums
array([1])