沿NumPy数组的轴计算唯一元素

Question

我有一个三维数组

A=np.array([[[1,1],
[1,0]],

[[1,2],
[1,0]],

[[1,0],
[0,0]]])

现在我想获得一个在给定位置具有非零值的数组,如果在该位置只出现唯一的非零值(或零).如果在该位置仅出现零或多于一个非零值,则它应该为零.对于上面的例子,我想

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

以来

• 在A[:,0,0]那里只有1s
• 在A[:,0,1]有0,1和2,所以一个以上的非零值
• 在A[:,1,0]有0和1,所以1被保留
• 在A[:,1,1]那里只有0s

我可以找到有多少非零元素np.count_nonzero(A, axis=0),但我想保留1s或2s,即使它们有几个.我看了一下,np.unique但它似乎不支持我想做的事情.

理想情况下,我喜欢一个函数np.count_unique(A, axis=0),它会以原始形状返回一个数组,例如[[1, 3],[2, 1]],所以我可以检查是否发生了3个或更多,然后忽略该位置.

---

所有我能想到的是一个列表理解迭代我想要获得的

[[len(np.unique(A[:, i, j])) for j in range(A.shape[2])] for i in range(A.shape[1])]

还有其他想法吗？

Answer 1

一种方法是使用A第一轴索引来沿其他两个轴设置相同长度的布尔数组，然后简单地计算沿其第一轴的非零值。可能有两种变体 - 一种保持不变，3D另一种是重塑2D为以获得一些性能优势，因为索引2D会更快。因此，这两个实现将是 -

def nunique_axis0_maskcount_app1(A):
    m,n = A.shape[1:]
    mask = np.zeros((A.max()+1,m,n),dtype=bool)
    mask[A,np.arange(m)[:,None],np.arange(n)] = 1
    return mask.sum(0)

def nunique_axis0_maskcount_app2(A):
    m,n = A.shape[1:]
    A.shape = (-1,m*n)
    maxn = A.max()+1
    N = A.shape[1]
    mask = np.zeros((maxn,N),dtype=bool)
    mask[A,np.arange(N)] = 1
    A.shape = (-1,m,n)
    return mask.sum(0).reshape(m,n)

运行时测试 -

In [154]: A = np.random.randint(0,100,(100,100,100))

# @B. M.'s soln
In [155]: %timeit f(A)
10 loops, best of 3: 28.3 ms per loop

# @B. M.'s soln using slicing : (B[1:] != B[:-1]).sum(0)+1
In [156]: %timeit f2(A)
10 loops, best of 3: 26.2 ms per loop

In [157]: %timeit nunique_axis0_maskcount_app1(A)
100 loops, best of 3: 12 ms per loop

In [158]: %timeit nunique_axis0_maskcount_app2(A)
100 loops, best of 3: 9.14 ms per loop

努巴法

使用nunique_axis0_maskcount_app2与使用 直接获取 C 级计数相同的策略numba，我们将得到 -

from numba import njit

@njit
def nunique_loopy_func(mask, N, A, p, count):
    for j in range(N):
        mask[:] = True
        mask[A[0,j]] = False
        c = 1
        for i in range(1,p):
            if mask[A[i,j]]:
                c += 1
            mask[A[i,j]] = False
        count[j] = c
    return count

def nunique_axis0_numba(A):
    p,m,n = A.shape
    A.shape = (-1,m*n)
    maxn = A.max()+1
    N = A.shape[1]
    mask = np.empty(maxn,dtype=bool)
    count = np.empty(N,dtype=int)
    out = nunique_loopy_func(mask, N, A, p, count).reshape(m,n)
    A.shape = (-1,m,n)
    return out

运行时测试 -

In [328]: np.random.seed(0)

In [329]: A = np.random.randint(0,100,(100,100,100))

In [330]: %timeit nunique_axis0_maskcount_app2(A)
100 loops, best of 3: 11.1 ms per loop

# @B.M.'s numba soln
In [331]: %timeit countunique2(A,A.max()+1)
100 loops, best of 3: 3.43 ms per loop

# Numba soln posted in this post
In [332]: %timeit nunique_axis0_numba(A)
100 loops, best of 3: 2.76 ms per loop