在两个不同维度的numpy数组之间获取所有唯一数字的最快方法

Question

我有两个不同维度的 numpy 数组，我需要在它们之间获取所有唯一的整数。

由于它们的维度不同，我决定为每个获取唯一的数字，将它们连接起来，然后获得唯一的值。

初步方法

import numpy as np
import time

array1 = np.random.randint(1, high=5000, size=(512, 4) , dtype=np.int32)
array2 = np.random.randint(1, high=5000, size=(512, 64) , dtype=np.int32)

start = time.time()
np.unique( np.concatenate( [np.unique(array1), np.unique(array2)] ) )
end = time.time()
print(end-start)

输出（平均 5 分）：

0.011

环境：谷歌 Colab

有没有更快的方法？

我的猜测是只执行一个 np.unique 操作，用不同大小的 numpy 数组来做这是一个挑战。

Answer 1

np.hstack工作速度会比您的初始方法略快（尽管差异不是很大）：

In [129]: %timeit np.unique(np.concatenate([np.unique(array1), np.unique(array2)]))                                      
1.74 ms ± 2.63 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [130]: %timeit np.unique(np.hstack((array1, array2)))                                                                 
1.58 ms ± 1.38 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 2

如果您的整数不会传播得太远，那么确实有一种更快的方法。 为简单起见，下面的代码假设非负整数小于 5000。

def use_mask():
    mask = np.zeros(5000,bool)
    mask[array1] = True
    mask[array2] = True
    return np.where(mask)

为了使它成为一个合适的函数，你必须找到 array1/2 组合的最小值和最大值，取差值，如果它太大，就回退到别的东西，否则从 array1/2 中减去最小值并或多或少地进行像我一样。