python - 按两列搜索大数组

Question

我有一个大数组，如下所示:

np.random.seed(42)

arr = np.random.permutation(np.array([
    (1,1,2,2,2,2,3,3,4,4,4),
    (8,9,3,4,7,9,1,9,3,4,50000)
]).T)

它没有排序，这个数组的行是唯一的，我也知道两列中值的界限，它们是 [0, n] 和 [0, k]。因此数组的最大可能大小为(n+1)*(k+1)，但实际大小更接近其对数。

我需要按两列搜索数组以找到 arr[row,:] = (i,j) 的 row，并返回 - 1 当数组中不存在 (i,j) 时。该函数的简单实现是:

def get(arr, i, j):
    cond = (arr[:,0] == i) & (arr[:,1] == j)
    if np.any(cond):
        return np.where(cond)[0][0]
    else:
        return -1

不幸的是，因为在我的例子中 arr 非常大(>90M 行)，所以效率非常低，特别是因为我需要调用 get( )多次。

或者，我尝试将其转换为带有 (i,j) 键的字典，这样

index[(i,j)] = row

可以通过以下方式访问:

def get(index, i, j):
   try:
      retuen index[(i,j)]
   except KeyError:
      return -1

这是可行的(并且在比我小的数据进行测试时要快得多)，但同样，通过即时创建字典

index = {}
for row in range(arr.shape[0]):
    i,j = arr[row, :]
    index[(i,j)] = row

就我而言，需要花费大量时间并消耗大量内存。我还考虑过首先对 arr 进行排序，然后使用 np.searchsorted 之类的东西，但这并没有引导我到任何地方。

所以我需要的是一个返回的快速函数get(arr, i, j)

>>> get(arr, 2, 3)
4
>>> get(arr, 4, 100)
-1 

Answer 1

部分解决方案是:

In [36]: arr
Out[36]: 
array([[    2,     9],
       [    1,     8],
       [    4,     4],
       [    4, 50000],
       [    2,     3],
       [    1,     9],
       [    4,     3],
       [    2,     7],
       [    3,     9],
       [    2,     4],
       [    3,     1]])

In [37]: (i,j) = (2, 3)

# we can use `assume_unique=True` which can speed up the calculation    
In [38]: np.all(np.isin(arr, [i,j], assume_unique=True), axis=1, keepdims=True)
Out[38]: 
array([[False],
       [False],
       [False],
       [False],
       [ True],
       [False],
       [False],
       [False],
       [False],
       [False],
       [False]])

# we can use `assume_unique=True` which can speed up the calculation
In [39]: mask = np.all(np.isin(arr, [i,j], assume_unique=True), axis=1, keepdims=True)

In [40]: np.argwhere(mask)
Out[40]: array([[4, 0]])

<小时/>

如果您需要最终结果作为标量，则不要使用 keepdims 参数并将数组转换为标量，例如:

    # we can use `assume_unique=True` which can speed up the calculation
In [41]: mask = np.all(np.isin(arr, [i,j], assume_unique=True), axis=1)

In [42]: np.argwhere(mask)
Out[42]: array([[4]])

In [43]: np.asscalar(np.argwhere(mask))
Out[43]: 4