方法 #1
这是一个使用NumPy向量化的方法,利用broadcasting
-
R,C = np.where((A[:,None,None] == B).any(-1))
out = np.split(C,np.flatnonzero(R[1:]>R[:-1])+1)
方法 #2
假设 A
和 B
是正数,我们可以将它们视为二维网格上的索引,其中B
可以被视为每行的列索引。一旦与 B
对应的二维网格就位,我们只需要考虑被 A
相交的列。最后,在这样的二维网格中获取 True
值的索引,即可得到 R
和 C
的值。这种方法应该更加节省内存。
因此,备选方法如下所示 -
ncols = B.max()+1
nrows = B.shape[0]
mask = np.zeros((nrows,ncols),dtype=bool)
mask[np.arange(nrows)[:,None],B] = 1
mask[:,~np.in1d(np.arange(mask.shape[1]),A)] = 0
R,C = np.where(mask.T)
out = np.split(C,np.flatnonzero(R[1:]>R[:-1])+1)
示例运行 -
In [43]: A
Out[43]: array([0, 1, 2, 3])
In [44]: B
Out[44]:
array([[3, 2, 0],
[0, 2, 1],
[2, 3, 1],
[3, 0, 1]])
In [45]: out
Out[45]: [array([0, 1, 3]), array([1, 2, 3]), array([0, 1, 2]), array([0, 2, 3])]
运行时间测试
将数据集大小扩大 100 倍
,以下是一个快速的运行时间测试结果 -
In [85]: def index_1din2d(A,B):
...: R,C = np.where((A[:,None,None] == B).any(-1))
...: out = np.split(C,np.flatnonzero(R[1:]>R[:-1])+1)
...: return out
...:
...: def index_1din2d_initbased(A,B):
...: ncols = B.max()+1
...: nrows = B.shape[0]
...: mask = np.zeros((nrows,ncols),dtype=bool)
...: mask[np.arange(nrows)[:,None],B] = 1
...: mask[:,~np.in1d(np.arange(mask.shape[1]),A)] = 0
...: R,C = np.where(mask.T)
...: out = np.split(C,np.flatnonzero(R[1:]>R[:-1])+1)
...: return out
...:
In [86]: A = np.unique(np.random.randint(0,10000,(400)))
...: B = np.random.randint(0,10000,(400,300))
...:
In [87]: %timeit [np.where((B == x).sum(axis = 1))[0] for x in A]
1 loop, best of 3: 161 ms per loop
In [88]: %timeit index_1din2d(A,B)
10 loops, best of 3: 91.5 ms per loop
In [89]: %timeit index_1din2d_initbased(A,B)
10 loops, best of 3: 33.4 ms per loop
进一步提升性能!
我们可以采用另一种方法以转置方式创建2D
网格。这个想法是为了避免在R,C = np.where(mask.T)
中进行转置,因为它似乎是瓶颈所在。因此,第二种方法的修改版本以及相关运行时间看起来像这样 -
In [135]: def index_1din2d_initbased_v2(A,B):
...: nrows = B.max()+1
...: ncols = B.shape[0]
...: mask = np.zeros((nrows,ncols),dtype=bool)
...: mask[B,np.arange(ncols)[:,None]] = 1
...: mask[~np.in1d(np.arange(mask.shape[0]),A)] = 0
...: R,C = np.where(mask)
...: out = np.split(C,np.flatnonzero(R[1:]>R[:-1])+1)
...: return out
...:
In [136]: A = np.unique(np.random.randint(0,10000,(400)))
...: B = np.random.randint(0,10000,(400,300))
...:
In [137]: %timeit index_1din2d_initbased(A,B)
10 loops, best of 3: 57.5 ms per loop
In [138]: %timeit index_1din2d_initbased_v2(A,B)
10 loops, best of 3: 25.9 ms per loop