索引超出范围的 Numpy 二维数组索引

Question

我在下面的代码中发现了一个实质性的瓶颈：

def get_value(matrix, index):
    if (index[0] >= 0 and index[1] >= 0 and
        index[0] < matrix.shape[0] and
        index[1] < matrix.shape[1]):
        return matrix[index[0], index[1]]
    return DEFAULT_VAL

给定一个二维矩阵和一个访问该矩阵的索引，它检查越界索引和 returns 给定索引处的值。否则，它 returns 一个 DEFAULT_VAL 在索引越界的情况下。

这个方法被调用了很多次（甚至上百万次），速度很慢。所以，我正在尝试使用 numpy 对其进行矢量化。不幸的是，我找不到办法。

如果我不必关心越界值，我会执行以下操作：

def get_values(matrix, indices):
    return matrix[indices[:,0], indices[:,1]]

我一直在想办法利用numpy来完成这个任务，但我还没有找到办法。

有没有办法做到这一点？

Answer 1

您显示的代码

def get_values(matrix, indices):
    return matrix[indices[:,0], indices[:,1]]

是你能做的最好的事情，因为 indices 是一个有两个值的元组。

您应该看看调用上述方法的最佳方式。我建议，如果可以，而不是用单个元组调用 get_values，调用可能有大量这样的元组。那么你至少可以尝试写一个矢量化版本的get_values。对于单个元组，您无法在此处进行矢量化。

矢量化方法

假设您的 indices 是一个大小为 n X 2 的 numpy 数组，其中 n 是索引的数量，而 2 对应于二维，那么您可以使用

index = np.random.randint(0,500, size=(10000,2))
matrix = np.random.randn(1000,1000)

def get_value(matrix, index, default_value=-1):
  result = np.zeros(len(index))+default_value
  mask = (index[:,0] < matrix.shape[0]) & (index[:,1] < matrix.shape[1])
  valid = index[mask]
  result[mask] = matrix[valid[:, 0], valid[:, 1]]
  return result


assert np.all(get_value(matrix, np.array(([0,1001],[1001,1001]))) == -1)

%timeit get_value(matrix, index, -1): 1 loop, best of 3: 264 ms per loop

索引超出范围的 Numpy 二维数组索引

Numpy 2D array indexing with indices out of bounds

python

numpy

optimization

matrix-indexing

矢量化方法