Python/NumPy：找到零的第一个索引，然后用零替换每一行之后的所有元素答案

【问题标题】：Python/NumPy: find the first index of zero, then replace all elements with zero after that for each rowPython/NumPy：找到零的第一个索引，然后用零替换每一行之后的所有元素
【发布时间】：2020-05-28 23:02:31
【问题描述】：

我有一个这样的 numpy 数组：

a = np.array([[1, 0, 1, 1, 1],
              [1, 1, 1, 1, 0],
              [1, 0, 0, 1, 1],
              [1, 0, 1, 0, 1]])

问题 1：如标题所示，我想在第一个零出现后将所有元素替换为零。结果应该是这样的：

a = np.array([[1, 0, 0, 0, 0],
              [1, 1, 1, 1, 0],
              [1, 0, 0, 0, 0],
              [1, 0, 0, 0, 0]])

问题2：如何像这个例子一样为每一行分割不同的列？因为我正在处理一个大尺寸的数组。如果有人能找到一种有效的方法来解决这个问题，请。非常感谢。

【问题讨论】：

标签： python numpy slice

【解决方案1】：

完成问题 1 的一种方法是使用numpy.cumprod

>>> np.cumprod(a, axis=1)
array([[1, 0, 0, 0, 0],
       [1, 1, 1, 1, 0],
       [1, 0, 0, 0, 0],
       [1, 0, 0, 0, 0]])

【讨论】：

不确定它的速度/效率，但它很漂亮，很有洞察力。

【解决方案2】：

问题 1：你可以像这样遍历数组：

for i in range(a.shape[0]):
    j = 0
    row = a[i]
    while row[j]>0:
        j += 1
    row[j+1:] = 0

这将就地更改数组。如果您对高性能感兴趣，the answers to this question 可能有助于更快地找到第一个零。 np.where 会为此扫描整个阵列，因此不是该任务的最佳选择。实际上，最快的解决方案将取决于数组条目的分布：如果那里有很多浮点数并且很少有零，则上面代码中的 while 循环平均会延迟中断，只需要写“几个”零。但是，如果您的样本数组中只有两个可能的条目，并且这些条目的发生概率相似（即~50%），那么将有很多零写入 a，以下会更快：

b = np.zeros(a.shape)
for i in range(a.shape[0]):
    j = 0
    a_row = a[i]
    b_row = b[i]
    while a_row[j]>0:
        b_row[j] = a_row[j]
        j += 1

问题 2：如果您的意思是按照处理某种第一次出现的类似标准对每一行单独切片，您可以简单地调整这种迭代模式。如果标准更具全局性（例如找到行的最大值），则存在像 np.where 这样的内置方法会更有效，但它可能会在一定程度上取决于标准本身，哪种选择是最好的。

【讨论】：

【解决方案3】：

问题 1：执行此操作的有效方法如下。

import numpy as np

a = np.array([[1, 0, 1, 1, 1],
              [1, 1, 1, 1, 0],
              [1, 0, 0, 1, 1],
              [1, 0, 1, 0, 1]])

for row in a:
    zeros = np.where(row == 0)[0]
    if (len(zeros)):# Check if zero exists
        row[zeros[0]:] = 0

print(a)

输出：

[[1 0 0 0 0]
 [1 1 1 1 0]
 [1 0 0 0 0]
 [1 0 0 0 0]]

问题 2：使用相同的数组，对于每一行 rowIdx，您可以有一个列数组 colIdxs 要从中提取。

rowIdx = 2
colIdxs = [1, 3, 4]
print(a[rowIdx, colIdxs])

输出：

[0 1 1]

【讨论】：

【解决方案4】：

我更喜欢 Ayrat 对第一个问题的创造性回答，但如果您需要为大尺寸的不同行分割不同的列，这可以帮助您：

indexer = tuple(np.s_[i:a.shape[1]] for i in (a==0).argmax(axis=1))
for i,j in enumerate(indexer):
    a[i,j]=0

索引器：

(slice(1, 5, None), slice(4, 5, None), slice(1, 5, None), slice(1, 5, None))

或：

indexer = (a==0).argmax(axis=1)
for i in range(a.shape[0]):
    a[i,indexer[i]:]=0

索引器：

[1 4 1 1]

输出：

[[1 0 0 0 0]
 [1 1 1 1 0]
 [1 0 0 0 0]
 [1 0 0 0 0]]

【讨论】：