获取嵌套列表中最大值的索引

Question

我有一个嵌套列表 a，里面有 N 个子列表，每个子列表里装着 M 个浮点数。我可以用 numpy 找到最大浮点数的索引，下面的 MWE 就是一个例子：

import numpy as np

def random_data(bot, top, N):
    # Generate some random data.
    return np.random.uniform(bot, top, N)

# Generate nested list a.
N, M = 10, 7  # number of sublists and length of each sublist
a = np.asarray([random_data(0., 1., M) for _ in range(N)])

# x,y indexes of max float in a.
print np.unravel_index(a.argmax(), a.shape)

需要注意的是，我把子列表的索引和浮点数的索引当作 x,y 坐标来用，其中 x 是子列表的索引，y 是该子列表里浮点数的索引。

现在我想找到最大浮点数的坐标/索引，但要加上一些限制条件。比如，我想在 x 的范围 [3:5] 和 y 的范围 [2:6] 内找到最大浮点数的 x,y 值。

这意味着我想在 a 里找最大浮点数，但只在 [3:5] 的子列表中查找，并且在这些子列表里只查找 [2:6] 范围内的浮点数。

我可以使用：

print np.unravel_index(a[3:5].argmax(), a[3:5].shape)

来限制 x 的范围，但返回的索引会因为列表被切片而 偏移，而且我想不出有什么办法可以这样得到 y 的索引。

Answer 1

你可以把'a'变成一个矩阵，这样就能方便地访问行和列了。接着，使用相同的基本命令可以获取限制后的索引。然后，你可以把限制索引的起始位置加到结果上，这样就能得到相对于整个矩阵的索引。

local_inds = np.unravel_index(A[3:5,2:6].argmax(),A[3:5,2:6].shape)
correct_inds = np.asarray(local_inds) + np.asarray([3,2])

如果你的索引限制比较复杂，这种方法就不太管用了。如果你有一份想要限制的x和y索引的列表，你可以这样做：

idx = [1,2,4,5,8]
idy = [0,3,6]
# make the subset of the data
b = np.asarray([[a[x][y] for y in idy] for x in idx])
# get the indices in the restricted data
indb = np.unravel_index(b.argmax(), b.shape)
# convert them back out
inda = (idx[indb[0]],idy[indb[1]])

Answer 2

另一种解决方案是将范围外的值设置为 np.inf：

import numpy as np

# Generate nested list a.
N, M = 10, 7  # number of sublists and length of each sublist
a = np.random.rand(N, M)

# x,y indexes of max float in a.
print np.unravel_index(a.argmax(), a.shape)

A = np.full_like(a, -np.inf)  # use np.inf if doing np.argmin
A[3:5, 2:6] = a[3:5, 2:6]
np.unravel_index(A.argmax(), A.shape)