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

编程问答 2022-05-22

问题描述

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

例如：

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

预期结果：

>> [[0,4],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：我也想找到左边的第一个非零元素。我想，找到正确点的解决方案可以很容易地找到左点。

解决方法

如果您的查询数组足够密集，您可以反转计算：找到一个与 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,[2,[1,1],[0,0]])
>>> np.stack((row,right_ind[row,3],4],3]])
>>> np.stack((row,col],1,3]])

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

,

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

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

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[:,arr=matrix)).T

结果是：

array([[0,3]],dtype=int64)

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

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

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

numpynumpy-ndarraypythonvectorization