通过第一列比较两个numpy数组，并拼接生成第三个numpy数组

Question

我有两个二维的numpy数组，用来绘制模拟结果。

这两个数组和的第一列都是时间间隔，第二列是要绘制的数据。这两个数组的形状不同，a是(500,2)，b是(600,2)。我想根据第一列来比较这两个numpy数组，并创建一个第三个数组，里面包含在的第一列中找到的匹配项。如果没有找到匹配，就在第三列添加0。

有没有什么numpy的小技巧可以做到这一点呢？

举个例子：

a=[[0.002,0.998],  
  [0.004,0.997],   
  [0.006,0.996],   
  [0.008,0.995],   
  [0.010,0.993]]   

b= [[0.002,0.666],  
    [0.004,0.665],  
    [0.0041,0.664], 
    [0.0042,0.664], 
    [0.0043,0.664], 
    [0.0044,0.663], 
    [0.0045,0.663], 
    [0.0005,0.663], 
    [0.006,0.663], 
    [0.0061,0.662],
    [0.008,0.661]] 

期望的输出结果

c= [[0.002,0.998,0.666],       
    [0.004,0.997,0.665],           
    [0.006,0.996,0.663],           
    [0.008,0.995,0.661],
    [0.010,0.993, 0   ]]  

Answer 1

下面的代码可以同时适用于numpy数组和普通的Python列表。

c = [[*x, y[1]] for x in a for y in b if x[0] == y[0]]
d = [[*x, 0] for x in a if x[0] not in [y[0] for y in b]]
c.extend(d)

要是有人胆子大一点，可以试着把这个代码写成一行。

Answer 2

import numpy as np
i = np.intersect1d(a[:,0], b[:,0])
overlap = np.vstack([i, a[np.in1d(a[:,0], i), 1], b[np.in1d(b[:,0], i), 1]]).T
underlap = np.setdiff1d(a[:,0], b[:,0])
underlap = np.vstack([underlap, a[np.in1d(a[:,0], underlap), 1], underlap*0]).T
fast_c = np.vstack([overlap, underlap])

这个方法是通过使用 intersect1d 来找出 a 和 b 的第一列的交集，然后用 in1d 来对比这个交集和第二列的数据。

vstack 是把输入的元素垂直堆叠起来，而转置操作是为了让数据的维度正确（这个操作非常快）。

接着，使用 setdiff1d 找出在 a 中但不在 b 中的时间，并在结果的第三列填上0。

最后会输出以下内容：

array([[ 0.002,  0.998,  0.666],
       [ 0.004,  0.997,  0.665],
       [ 0.006,  0.996,  0.   ],
       [ 0.008,  0.995,  0.   ],
       [ 0.01 ,  0.993,  0.   ]])

Answer 3

我很快就想到了一个解决办法：

import numpy as np

a = np.array([[0.002, 0.998],
     [0.004, 0.997],
     [0.006, 0.996],
     [0.008, 0.995],
     [0.010, 0.993]])

b = np.array([[0.002, 0.666],
     [0.004, 0.665],
     [0.0041, 0.664],
     [0.0042, 0.664],
     [0.0043, 0.664],
     [0.0044, 0.663],
     [0.0045, 0.663],
     [0.0005, 0.663],
     [0.0006, 0.663],
     [0.00061, 0.662],
     [0.0008, 0.661]])


c = []
for row in a:
    index = np.where(b[:,0] == row[0])[0]
    if np.size(index) != 0:
      c.append([row[0], row[1], b[index[0], 1]])
    else:
      c.append([row[0], row[1], 0])

print c

正如上面的评论所提到的，似乎存在一个数据录入错误。