numpy 中集合的向量化相对补集

Vectorized relative complement of sets in numpy

我有 np.arange(n) A 和一个 numpy 数组 B 的非相交子数组 - 除法初始数组转化为k个连续数数组
一个例子是:

A = [0, 1, 2, 3, 4, 5, 6]
B = [[0, 1], [2, 3, 4], [5, 6]]

对于 B 的每个子数组 C 我必须计算 A\C (其中 \ 是对集合的操作,因此结果是 A 中不在 B 中的所有元素的 numpy 数组。 我当前的解决方案达到了时间限制:

import numpy as np
for C in B:
    ans.append(np.setdiff1d(A, C))
return ans

我想通过矢量化来加速它,但我不知道该怎么做。我试图删除循环,只留下像 setxor1d 和 setdiff1d 这样的函数,但失败了。

我假设 A 和 B 的子数组已排序并且具有唯一元素。然后对于我下面的示例,将 10**6 个整数分成由以下代码生成的 100 个子数组。

np.random.seed(0)
A = np.sort(np.unique(np.random.randint(0,10**10,10**6)))
B = np.split(A, np.sort(np.random.randint(0,10**6-1,99)))

您可以通过设置 unique=True 将时间减半。并通过仅对 A 中位于 B 的特定子集中最大和最小数字之间的数字执行 setminus in 来将时间缩短 3 倍。我意识到我的示例是最佳情况优化以提供帮助,所以我不确定这对您的真实世界示例有何影响。你将不得不尝试。

boundaries = [x[i] for x in B for i in [0,-1]]
boundary_idx = np.searchsorted(A, boundaries).reshape(-1,2)
[np.concatenate([A[:x[0]], 
                 np.setdiff1d(A[x[0]:x[1]+1], b, assume_unique=True), 
                 A[x[1]+1:]])
         for b,x in zip(B, boundary_idx)]