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

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 语句转换为这个二维数组...

Answer 1

您已经有了很好的答案，但我想使用函数 np.maximum.accumulate 提出一个可能更快的变体.由于您的一维数组方法使用 max/where ，您可能还会发现这种方法非常直观。 (编辑:下面添加了更快的 Cython 实现)。

整体方法与其他方法非常相似；掩码是通过以下方式创建的:

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

这行代码做了以下事情:

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

• np.maximum.accumulate用于返回沿每一行的累积最大值(即 axis=1 )。例如 [False, True, False]将变为 [False, True, True] .

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

那么剩下要做的就是使用 bool 数组作为掩码来设置True值为零。

---

借用 @Divakar's answer 中的计时方法和两个函数，以下是我提出的方法的基准:

# 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 points out ，每个 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