获取二维数组每列最后一个负数的索引

Question

我正在尝试获取每列数组中最后一个负值的索引（以便在其后进行切片）。 一个在一维向量上的简单工作示例是：

import numpy as np

A = np.arange(10) - 5
A[2] = 2
print A # [-5 -4  2 -2 -1  0  1  2  3  4]

idx = np.max(np.where(A <= 0)[0])
print idx # 5

A[:idx] = 0
print A # [0 0 0 0 0 0 1 2 3 4]

现在我想在二维数组的每一列上执行相同的操作：

A = np.arange(10) - 5
A[2] = 2
A2 = np.tile(A, 3).reshape((3, 10)) - np.array([0, 2, -1]).reshape((3, 1))
print A2
# [[-5 -4  2 -2 -1  0  1  2  3  4]
#  [-7 -6  0 -4 -3 -2 -1  0  1  2]
#  [-4 -3  3 -1  0  1  2  3  4  5]]

我希望获取：

print A2
# [[0 0 0 0 0 0 1 2 3 4]
#  [0 0 0 0 0 0 0 0 1 2]
#  [0 0 0 0 0 1 2 3 4 5]]

但我无法想出如何将max/where语句翻译成这个二维数组...

- Thomas Leonard

---

是的，我需要... [-3,-4,-5,3,4,-7,8] 必须给我 [0,0,0,0,0,0,8]。 - Thomas Leonard

针对您的实际用例：典型的数组大小是多少？负数的数量和分布通常是怎样的？ - tom10

最终，每行的多少部分将被清零？ - tom10

很难说...可能是行的80%... - Thomas Leonard

好的。根据这些数字，我认为迭代行会更快。我会尽快发布一个示例（可能需要几天时间）。另一方面，如果您对现有内容感到满意，请告诉我，我会放弃修改。 - tom10

Answer 1

你已经有了很好的答案，但我想提出一个可能更快的变化，使用函数np.maximum.accumulate。由于你在处理一维数组时使用了max/where，所以你可能会发现这种方法非常直观。（编辑：下面添加了更快的Cython实现）。

总体方法与其他方法非常相似；掩码是用以下方式创建的：

np.maximum.accumulate((A2 < 0)[:, ::-1], axis=1)[:, ::-1]

这行代码的作用如下：

• (A2 < 0) 创建一个布尔数组，指示值是否为负。索引 [:, ::-1] 将其从左到右翻转。

• 使用 np.maximum.accumulate 返回沿每行（即 axis=1）的累积最大值。例如，[False, True, False] 变成了 [False, True, True]。

• 最后的索引操作 [:, ::-1] 将这个新的布尔数组从左到右翻转。

然后只需要将布尔数组用作掩码，将 True 值设置为零即可。

---

借鉴自@Divakar 的答案的计时方法和两个函数，以下是我提出的方法的基准测试结果：

# method using np.maximum.accumulate
def accumulate_based(A2):
    A2[np.maximum.accumulate((A2 < 0)[:, ::-1], axis=1)[:, ::-1]] = 0
    return A2

# large sample array
A2 = np.random.randint(-4, 10, size=(100000, 100))
A2c = A2.copy()
A2c2 = A2.copy()

时间为：

In [47]: %timeit broadcasting_based(A2)
10 loops, best of 3: 61.7 ms per loop

In [48]: %timeit cumsum_based(A2c)
10 loops, best of 3: 127 ms per loop

In [49]: %timeit accumulate_based(A2c2) # quickest
10 loops, best of 3: 43.2 ms per loop

因此，对于这种大小和形状的数组，使用np.maximum.accumulate可能比下一个最快的解决方案快多达30％。

---

正如@tom10指出的那样，每个NumPy操作都会处理整个数组，在需要多次操作才能得到结果时可能效率不高。一种迭代方法可以一次遍历数组就完成计算，这可能更有效。

下面是一个用Cython编写的朴素函数，其速度可能比纯NumPy方法快两倍以上。

可以尝试使用memory views进一步加快此函数的速度。

cimport cython
import numpy as np
cimport numpy as np

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def cython_based(np.ndarray[long, ndim=2, mode="c"] array):
    cdef int rows, cols, i, j, seen_neg
    rows = array.shape[0]
    cols = array.shape[1]
    for i in range(rows):
        seen_neg = 0
        for j in range(cols-1, -1, -1):
            if seen_neg or array[i, j] < 0:
                seen_neg = 1
                array[i, j] = 0
    return array

该函数从每一行的末尾开始向前遍历，在遇到负数后开始将其后面的值设为零。 测试它是否有效：

A2 = np.random.randint(-4, 10, size=(100000, 100))
A2c = A2.copy()

np.array_equal(accumulate_based(A2), cython_based(A2c))
# True

比较函数的性能:

In [52]: %timeit accumulate_based(A2)
10 loops, best of 3: 49.8 ms per loop

In [53]: %timeit cython_based(A2c)
100 loops, best of 3: 18.6 ms per loop

- Alex Riley

Answer 2

假设您想将每行的所有元素设置为零，直到最后一个负元素（根据问题中列出的样例输出），这里可以提供两种方法。 方法一 这种方法基于np.cumsum生成要设置为零的元素掩码，如下所示 -

# Get boolean mask with TRUEs for each row starting at the first element and 
# ending at the last negative element
mask = (np.cumsum(A2[:,::-1]<0,1)>0)[:,::-1]

# Use mask to set all such al TRUEs to zeros as per the expected output in OP 
A2[mask] = 0

样例运行 -

In [280]: A2 = np.random.randint(-4,10,(6,7)) # Random input 2D array

In [281]: A2
Out[281]: 
array([[-2,  9,  8, -3,  2,  0,  5],
       [-1,  9,  5,  1, -3, -3, -2],
       [ 3, -3,  3,  5,  5,  2,  9],
       [ 4,  6, -1,  6,  1,  2,  2],
       [ 4,  4,  6, -3,  7, -3, -3],
       [ 0,  2, -2, -3,  9,  4,  3]])

In [282]: A2[(np.cumsum(A2[:,::-1]<0,1)>0)[:,::-1]] = 0 # Use mask to set zeros

In [283]: A2
Out[283]: 
array([[0, 0, 0, 0, 2, 0, 5],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 3, 5, 5, 2, 9],
       [0, 0, 0, 6, 1, 2, 2],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 9, 4, 3]])

方法二

这个方法起始于从@tom10的答案中找到最后一个负元素的索引，并发展成使用broadcasting的掩码查找方法，以获得与方法一类似的所需输出。

# Find last negative index for each row
last_idx = A2.shape[1] - 1 - np.argmax(A2[:,::-1]<0, axis=1)

# Find the invalid indices (rows with no negative indices)
invalid_idx = A2[np.arange(A2.shape[0]),last_idx]>=0

# Set the indices for invalid ones to "-1"
last_idx[invalid_idx] = -1

# Boolean mask with each row starting with TRUE as the first element 
# and ending at the last negative element
mask = np.arange(A2.shape[1]) < (last_idx[:,None] + 1)

# Set masked elements to zeros, for the desired output
A2[mask] = 0

---

运行时测试 -

函数定义：

def broadcasting_based(A2):
    last_idx = A2.shape[1] - 1 - np.argmax(A2[:,::-1]<0, axis=1)
    last_idx[A2[np.arange(A2.shape[0]),last_idx]>=0] = -1
    A2[np.arange(A2.shape[1]) < (last_idx[:,None] + 1)] = 0
    return A2

def cumsum_based(A2):    
    A2[(np.cumsum(A2[:,::-1]<0,1)>0)[:,::-1]] = 0    
    return A2

运行时：

In [379]: A2 = np.random.randint(-4,10,(100000,100))
     ...: A2c = A2.copy()
     ...: 

In [380]: %timeit broadcasting_based(A2)
10 loops, best of 3: 106 ms per loop

In [381]: %timeit cumsum_based(A2c)
1 loops, best of 3: 167 ms per loop

验证结果 -

In [384]: A2 = np.random.randint(-4,10,(100000,100))
     ...: A2c = A2.copy()
     ...: 

In [385]: np.array_equal(broadcasting_based(A2),cumsum_based(A2c))
Out[385]: True

- Divakar

Answer 3

通常找到第一个比找到最后一个更容易更快，因此在这里我反转了数组，然后找到第一个负数（使用OP的版本A2）：

im = A2.shape[1] - 1 - np.argmax(A2[:,::-1]<0, axis=1)

# [4 6 3]      # which are the indices of the last negative in A2

此外，如果您有许多负数的大型数组，使用非numpy方法可能会更快，因为您可以短路搜索。也就是说，numpy将对整个数组进行计算，因此如果您有10000个元素，但通常会在前10个元素中（通过反向搜索）找到负数，则纯Python方法可能会更快。

总体而言，迭代行在后续操作中可能更快。例如，如果下一步是乘法，则只需乘以非零末尾的切片，或者找到最长的非零部分并处理截断的数组可能更快。

这基本上取决于每行的负数数量。如果每行有1000个负数，则平均而言，非零段将是您完整行长度的1/1000，因此您可以通过查看末尾来获得1000倍的加速。问题中给出的简短示例非常适合理解和回答基本问题，但是当您的最终应用程序是非常不同的用例时，我不会认真考虑时间测试；特别是因为使用迭代的分数时间节省与数组大小成比例改善（假设负数具有恒定比率和随机分布）。

- tom10

Answer 4

您可以访问单独的行：

A2[0] == array([-5, -4,  2, -2, -1,  0,  1,  2,  3,  4])

- kanu

---

我知道我可以循环每一行并做相同的事情，但是在我的实际情况中，我的数组包含数百万行，因此我需要一些高效的方法，意味着不使用循环。 - Thomas Leonard

@thomleo，将这些信息包含在你的问题和/或标题中可能是一个好主意，这对回答者和未来的读者都有帮助。 - user