在python中查找给定元素右侧的第一个非零元素的索引

Question

我有一个 2D numpy.ndarray。给定一个位置列表，我想找到同一行中给定元素右侧的第一个非零元素的位置。是否可以将其矢量化？我有一个庞大的数组，循环花费了太多时间。

例如：

matrix = numpy.array([
    [1, 0, 0, 1, 1], 
    [1, 1, 0, 0, 1], 
    [1, 0, 0, 0, 1], 
    [1, 1, 1, 1, 1], 
    [1, 0, 0, 0, 1]
])
query = numpy.array([[0,2], [2,1], [1,3], [0,1]])

预期结果：

>> [[0,3], [2,4], [1,4], [0,3]]

目前我正在使用 for 循环执行此操作，如下所示

for query_point in query:
    y, x = query_point
    result_point = numpy.min(numpy.argwhere(self.matrix[y, x + 1:] == 1)) + x + 1
    print(f'{y}, {result_point}')

PS：我也想找到左边的第一个非零元素。我想，找到右点的解决方案可以很容易地找到左点。

Answer 1

我想出了一个解决方案来获得两个您想要的索引， 即从指示位置向左和向右。

首先定义如下函数，获取行号和两个索引：

def inds(r, c, arr):
    ind = np.nonzero(arr[r])[0]
    indSlice = ind[ind < c]
    iLeft = indSlice[-1] if indSlice.size > 0 else None
    indSlice = ind[ind > c]
    iRight = indSlice[0] if indSlice.size > 0 else None
    return r, iLeft, iRight

参数：

• r 和 c 是行号（在源数组中）和“开始” 此行中的索引，
• arr 是要查看的数组（matrix 将在此处传递）。

然后定义这个函数的矢量化版本：

indsVec = np.vectorize(inds, excluded=['arr'])

为了得到结果，运行：

result = np.vstack(indsVec(query[:, 0], query[:, 1], arr=matrix)).T

结果是：

array([[0, 0, 3],
       [2, 0, 4],
       [1, 1, 4],
       [0, 0, 3]], dtype=int64)

您的预期结果是左右列（行号 以及“起始”位置之后的第一个非零元素的索引。

中间列是“开始”位置之前最后一个非零元素的索引。

此解决方案可以抵抗“不存在”的情况（如果没有 任何“之前”或“之后”非零元素）。在这种情况下，相应的 索引返回为无。

Answer 2

如果您的查询数组足够密集，您可以反转计算：找到一个与matrix 大小相同的数组，它为每个位置提供同一行中下一个非零元素的索引。那么你的问题就变成了将query 应用于这个索引数组的问题之一，numpy 直接支持。

实际上找到左索引要容易得多，所以让我们从它开始。我们可以将matrix 转换成这样的索引数组：

r, c = np.nonzero(matrix)
left_ind = np.zeros(matrix.shape, dtype=int)
left_ind[r, c] = c

现在您可以使用 np.maximum 找到前一个非零元素的索引，类似于此答案中的方法：https://stackoverflow.com/a/48252024/2988730：

np.maximum.accumulate(left_ind, axis=1, out=left_ind)

现在您可以直接索引到ind 以获取之前的非零列索引：

left_ind[query[:, 0], query[:, 1]]

或

left_ind[tuple(query.T)]

现在要对正确的索引做同样的事情，您需要反转数组。但是随后您的索引不再上升，并且您可能会覆盖第一列中的任何零。为了解决这个问题，除了反转数组之外，您还需要反转索引的顺序：

right_ind = np.zeros(matrix.shape, dtype=int)
right_ind[r, c] = matrix.shape[1] - c

您也可以使用大于matrix.shape[1] 的任何数字作为常量。重要的是反向索引都大于零，因此np.maximum.accumulate 会覆盖零。现在您可以在反转数组上以相同的方式使用np.maximum.accumulate：

right_ind = matrix.shape[1] - np.maximum.accumulate(right_ind[:, ::-1], axis=1)[:, ::-1]

在这种情况下，我建议不要使用out=right_ind，因为right_ind[:, ::-1] 是同一个缓冲区的视图。该操作是缓冲的，但如果您的行大小足够大，您可能会无意中覆盖数据。

现在你可以像以前一样索引数组了：

right_ind[(*query.T,)]

在这两种情况下，您都需要与query 的第一列堆叠，因为那是行键：

>>> row, col = query.T
>>> np.stack((row, left_ind[row, col]), -1)
array([[0, 0],
       [2, 0],
       [1, 1],
       [0, 0]])
>>> np.stack((row, right_ind[row, col]), -1)
array([[0, 3],
       [2, 4],
       [1, 4],
       [0, 3]])
>>> np.stack((row, left_ind[row, col], right_ind[row, col]), -1)
array([[0, 0, 3],
       [2, 0, 4],
       [1, 1, 4],
       [0, 0, 3]])

如果您计划一次或在整个程序中对数组中的大部分行进行采样，这将帮助您加快速度。另一方面，如果您只需要访问一个小子集，则可以将此技术仅应用于您需要的行。