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

Question

我有一个二维 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

如果你的查询数组足够密集，你可以反向计算:找到一个与 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]是同一个缓冲区的 View 。该操作是缓冲的，但如果您的行大小足够大，您可能会无意中覆盖数据。
现在您可以像以前一样对数组进行索引:

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]])

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