如何高效提取numpy数组中匹配指定ID的行（无循环、保序）

我有两个NumPy数组：

• 单列数组points，存储了数千万个点的ID；
• 三列数组coordinates，存储了数千万个点的ID、x坐标、y坐标。

数组结构如下：

import numpy as np

points = np.array([[id], [id], ..., [id]])  # 点ID的单列数组
coordinates = np.array([[id, x, y], ..., [id, x, y]]) # 包含ID与对应坐标的三列数组

我需要实现类似SQL语句SELECT * FROM coordinates WHERE id IN points的效果，提取coordinates中ID存在于points中的所有行，且结果顺序要与points一致。

约束条件 #

• 两个数组的ID排序方式不同
• 不能使用索引循环，需实现高效算法

我之前试过循环实现，但效率太低想规避：

indexes = []
for i in range(len(points)):
    indexes.append(np.where(coordinates[:,0] == points[i])[0][0])
result = coordinates[indexes, :]

针对千万级数据量，推荐使用基于排序+二分查找的向量化操作，完全规避循环，时间复杂度为O(n log n)，远优于循环实现的O(n²)。

代码实现 #

import numpy as np

# 将points从单列数组转为一维，避免维度匹配问题
points_flat = points.flatten()

# 提取coordinates的ID列，生成排序后的索引与排序后的ID数组
coord_ids = coordinates[:, 0]
sorted_indices = np.argsort(coord_ids)
sorted_coord_ids = coord_ids[sorted_indices]

# 用二分查找快速定位每个points ID在排序后ID数组中的位置
positions = np.searchsorted(sorted_coord_ids, points_flat)

# 可选：验证匹配结果，过滤points中不存在于coordinates的ID（根据数据情况选择）
mask = sorted_coord_ids[positions] == points_flat
valid_positions = positions[mask]

# 根据定位结果提取对应行，保证顺序与points一致
result = coordinates[sorted_indices[valid_positions]]

# 若确认points中所有ID都存在于coordinates，可直接用以下代码：
# result = coordinates[sorted_indices[positions]]

方案优势 #

• 所有操作均为NumPy底层优化的向量化实现，用C语言执行，比Python循环效率提升数个数量级；
• 仅需一次排序（O(m log m)，m为coordinates长度），后续每个ID查找仅需O(log m)时间，整体复杂度远低于循环的O(n*m)；
• 完美适配千万级别的大数据量处理需求，同时保证结果顺序与points完全一致。

内容的提问来源于stack exchange，提问作者Coquelicot