首页 > 解决方案 > 从两个具有接近值的二维数组中查找行索引

问题描述

我有一个数组a

a = np.array([[4, 4],
              [5, 4],
              [6, 4],
              [4, 5],
              [5, 5],
              [6, 5],              
              [4, 6],
              [5, 6],
              [6, 6]])

和一个数组b

b = np.array([[4.001,4], 
              [8.001,4],
              [5,4.0003],
              [5.9999,5]])

我想找到它的索引a值非常接近b. 如果b数组的值与 I 中的值完全相同,则a可以使用以下代码。

np.where((a==b[:,None]).all(-1))[1]

为了清楚起见;我希望代码返回以下内容:[0,1,5] 这些是 a 的索引,非常接近 b 中的行。带有值的 inb[8.001,4]被丢弃,因为它不在a. 我认为将上面的代码与 np.allclose() 结合可以解决它,但是我似乎无法弄清楚如何做到这一点,你能帮我吗?

标签: pythonarraysnumpy

解决方案


而不是==,使用np.linalg.normon来获取 in 中的每一行到每一行a - b[:, np.newaxis]的距离。ab

如果a并且b有很多行,这将使用大量内存:例如,如果每个行有 10,000 行,则vecdiff下面的数组将是 10,000×10,000 或 100,000,000 个元素;使用双打这是 800MB。

In [50]: a = np.array([[4, 4], 
    ...:               [5, 4], 
    ...:               [6, 4], 
    ...:               [4, 5], 
    ...:               [5, 5], 
    ...:               [6, 5],               
    ...:               [4, 6], 
    ...:               [5, 6], 
    ...:               [6, 6]]) 
    ...:                                                                        

In [51]: b = np.array([[4.001,4],  
    ...:               [8.001,4], 
    ...:               [5,4.0003], 
    ...:               [5.9999,5]])                                             

In [52]: vecdist = np.linalg.norm(a - b[:, np.newaxis], axis=-1)                

In [53]: closeidx = np.flatnonzero(vecdist.min(axis=0) < 1e-2)                  

In [54]: print(closeidx)                                                        
[0 1 5]

推荐阅读