如何通过矩阵numpy向量化循环

Question

假设我有一个100000 x 100的矩阵

import numpy as np

mat = np.random.randint(2, size=(100000,100))

我希望遍历此矩阵，并且如果每个矩阵都row完全包含1或0，则希望将state变量更改为该值。如果状态未更改，则希望设置的整个row值state。的初始值为state0。

天真的for循环可以做到如下

state = 0

for row in mat:
    if set(row) == {1}:
        state = 1
    elif set(row) == {0}:
        state = 0
    else:
        row[:] = state

然而，当矩阵的大小增加时，这花费了不切实际的时间。有人numpy可以向我指出如何利用该循环矢量化并加快速度的方向吗？

所以对于样本输入

array([[0, 1, 0],
       [0, 0, 1],
       [1, 1, 1],
       [0, 0, 1],
       [0, 0, 1]])

在这种情况下，预期输出为

array([[0, 0, 0],
       [0, 0, 0],
       [1, 1, 1],
       [1, 1, 1],
       [1, 1, 1]])

Answer 1

这是一个简单快速的 numpy 方法：

import numpy as np

def pp():
    m,n = a.shape
    A = a.sum(axis=1)    
    A = np.where((A==0)|(A==n))[0]
    if not A.size:
        return np.ones_like(a) if state else np.zeros_like(a)
    st = np.concatenate([np.arange(A[0]!=0), A, [m]])
    v = a[st[:-1],0]
    if A[0]:
        v[0] = state
    return np.broadcast_to(v.repeat(st[1:]-st[:-1])[:,None],(m,n))

我用这个做了一些计时

state=0
a = (np.random.random((100000,100))<np.random.random((100000,1))).astype(int)

简单的测试用例：

0.8655898020006134   # me
4.089095343002555    # Alain T.
2.2958932030014694   # Divakar 1
2.2178015549980046   # & 2