有没有np.where（）更快的替代方法？

Question

我有一组100个数据文件，其中包含有关粒子的信息（ID，速度，位置等）。我需要从每个粒子中选择10000个具有特定ID号的特定粒子。我的操作方式如下

AddJob

列表“ chosen_id_array”包含所有此类ID。数据文件是根据列表索引进行结构化的。

由于某些原因，此代码段运行非常缓慢，我正在寻找一种更快，更有效的替代方法。提前非常感谢您。 ：）

Answer 1

使用字典，您可以存储归因于粒子ID的位置信息，并利用O（1）对字典进行查找缩放：

# What the data in a single file would look like:
data = {1:[0.5,0.1,1.], 4:[0.4,-0.2,0.1], ...}
# A lookup becomes very simple syntactically:
for ID in chosen_id_arr:
    x, y, z = data[ID]
    # Here you can process the obtained x,y,z.

这比numpy查找要快得多。关于循环中位置数据的处理，您可以考虑为不同的粒子ID具有单独的位置列表，但我认为这不在问题范围内。熊猫包也可以在那里提供帮助。