来自较大 numpy 数组的唯一 numpy 数组（非元素）的行索引答案

【问题标题】：Row indices of unique numpy array (not element) from a larger numpy array来自较大 numpy 数组的唯一 numpy 数组（非元素）的行索引
【发布时间】：2022-01-22 12:33:25
【问题描述】：

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

'a' 是形状 [30000,2] 并包含成对的 x,y 坐标。 'b' 的形状为 [10,000,000,3] 并包含 x,y,z 坐标。

来自'a' 的x,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!!!)

我确信这是一个简单的问题，但它让我很难过。

实现这一目标的最有效方法是什么？（尤其是对于较大的数据集）

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

您可以将二维数组转换为一维视图（参见this answer），然后使用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])

【讨论】：

【解决方案2】：

如果您不想展平数组，另一种选择是使它们具有相同的大小并按元素进行比较：

a, b = a.T, b.T
tile_a = np.tile(a, b[: , :2].shape[0]).reshape(a.shape[0] * b[: , :2].shape[0], a.shape[1])
indices = np.argwhere((tile_a == np.concatenate([b[:, :2]]*a.shape[0])).all(axis=1))
indices[indices > 0] -= 1 

print(np.squeeze(b[indices // a.shape[0]], axis=1)[:, 2])
#[11 13 15 18]

【讨论】：

回答有用吗？ @DavidMegson-Smith