首页 > 解决方案 > 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 之类的函数,但失败了。

标签: pythonnumpyoptimizationsetvectorization

解决方案


我假设 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。并通过仅对位于 B 的特定子集中的最大和最小数字之间的 A 中的数字进行 setminus 将时间缩短 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)]

推荐阅读