较大numpy数组中唯一numpy数组(而非元素)的行索引



我有两个numpy数组'a'和'b'。

"a"的形状为[30000,2],包含x,y坐标对b'的形状为[10000000,3],包含x、y、z坐标。

x、 "a"中的y坐标对在"b"中总是只出现一次。我想有效地从"b"中提取相应的z坐标。

这里有一个简单的例子。。。

a = np.array([[1,2], [3,4], [5,6], [8,9]]).T
b = np.array([[1,2,11], [1,3,12], [3,4,13], [4,5,14],[5,6,15], [6,7,16], [7,8,17], [8,9,18]]).T 

将返回[0,2,4,7]的行索引,使得z = [11, 13, 15, 18]

显然,这可以通过2个for循环(YUCK!!(来实现

我确信这是一个简单的问题,但它让我不知所措。

实现这一目标最有效的方法是什么?(特别是对于较大的数据集(

您可以将2D阵列转换为1D视图(请参阅此答案(,然后使用numpy.isin:

def view1D(a, b):
a = np.ascontiguousarray(a)
b = np.ascontiguousarray(b)
void_dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
return a.view(void_dt).ravel(),  b.view(void_dt).ravel()
A,B = view1D(a.T, b[:2].T)
b.T[np.isin(B, A)][:,2]
# array([11, 13, 15, 18])

最新更新