Numpy - 3D 数组中重复行的计数

Question

我希望计算 3D NumPy 数组中唯一行的数量。取以下数组：

a = np.array([[[1, 2], [1, 2], [2, 3]], [[2, 3], [2, 3], [3, 4]], [[1, 2], [1, 2], [1, 2]]])

我想要的输出是一个与 3-D 数组的轴 0 长度相同的 1-D 数组。 array([2, 2, 1]).

在此示例中，输出将为 2, 2, 1，因为在第一个分组中 [1, 2] 和 [2, 3] 是唯一值，在第二个分组中 [2, 3] 和 [3, 4] 是唯一值，第三组 [1, 2] 是“唯一”值。也许我在这种情况下错误地使用了 unique ，但这是我想要计算的。

我遇到的困难是唯一行的数量会有所不同。如果我使用np.unique，结果广播如下图：

>>> np.unique(a, axis=1)
array([[[1, 2],
        [2, 3]],

       [[2, 3],
        [3, 4]],

       [[1, 2],
        [1, 2]]])

我知道我可以遍历每个二维数组并使用np.apply_along_axis，如this answer 中所述。

但是，我正在处理像 (1 000 000, 256, 2) 这样大的数组，所以如果可能的话，我宁愿避免循环。

Answer 1

为每个 2D 计划调用 np.unique 似乎非常慢。实际上，它是 np.unique，它很慢，并不是真正的纯 Python 循环。

更好的方法是手动使用 Numba（使用 dict）。虽然这种策略更快，但它不是灵丹妙药。但是，尽管dict 的访问速度不是很快，但此实现可以很容易地并行化以显着加快运行速度。这是实现：

import numpy as np
import numba as nb

@nb.njit('i4[::1](i4[:,:,::1])', parallel=True)
def compute_unique_count(data):
    n,m,o = data.shape
    assert o == 2
    res = np.empty(n, dtype=np.int32)
    for i in nb.prange(n):
        tmp = dict()
        for j in range(m):
            tmp[(data[i, j, 0], data[i, j, 1])] = True
        res[i] = len(tmp)
    return res

由于数字是小有界整数，因此有一种有效的方法可以使计算速度更快。实际上，关联表可用于标记已找到的值，其中数组的索引是dict 的键，而数组的值定义了迄今为止找到的键。为了性能，可以使用表达式id = point[0] * maxi + point[1] 将数组展平，其中maxi 是界限（假设所有值都严格小于它）。为了性能，避免了分支，maxi 应该是 2 的幂，通常 >= 64（由于缓存行冲突、缓存抖动和错误共享等低级缓存影响）。结果实现非常快。

@nb.njit('int32[::1](int32[:,:,::1])', parallel=True)
def compute_unique_count_fastest(data):
    n,m,o = data.shape
    assert o == 2
    maxi = 64
    threadCount = nb.get_num_threads()
    res = np.empty(n, dtype=np.int32)
    globalUniqueVals = np.zeros((threadCount, maxi * maxi), dtype=np.uint8)
    for i in nb.prange(n):
        threadId = nb.np.ufunc.parallel._get_thread_id()
        uniqueVals = globalUniqueVals[threadId]
        uniqueVals.fill(0) # Reset the associative table
        uniqueCount = 0
        for j in range(m):
            idx = data[i, j, 0] * maxi + data[i, j, 1]
            uniqueCount += uniqueVals[idx] == 0
            uniqueVals[idx] = 1
        res[i] = uniqueCount
    return res

这是我的机器（i5-9600KF，6 核）上的时序，数组大小为 (1_000_000, 256, 2)，包含从 0 到 40 的随机 32 位整数：

np.unique in a comprehension list:  78 000 ms
compute_unique_count:                2 050 ms 
compute_unique_count_fastest:           57 ms

最后一个实现比原始实现快 1370 倍。