python - 使用标识符列对数组行进行排序以匹配另一个数组的顺序

Question

我有两个这样的数组:

A = [[111, ...],          B = [[222, ...],
     [222, ...],               [111, ...],
     [333, ...],               [333, ...],
     [555, ...]]               [444, ...],
                               [555, ...]]

其中第一列包含标识符，其余列是一些数据，其中 B 的列数远大于 A 的列数。标识符是唯一的。 A 中的行数可以少于 B 中的行数，因此在某些情况下需要空的间隔行。
我正在寻找一种有效的方法来将矩阵 A 的行与矩阵 B 相匹配，以便结果看起来像这样:

A = [[222, ...],
     [111, ...],
     [333, ...],
     [nan, nan], #could be any unused value
     [555, ...]]

我可以对两个矩阵进行排序或编写一个 for 循环，但这两种方法看起来都很笨拙......有更好的实现吗？

Answer 1

这是使用 np.searchsorted 的矢量化方法-

# Store the sorted indices of A
sidx = A[:,0].argsort()

# Find the indices of col-0 of B in col-0 of sorted A
l_idx = np.searchsorted(A[:,0],B[:,0],sorter = sidx)

# Create a mask corresponding to all those indices that indicates which indices
# corresponding to B's col-0 match up with A's col-0
valid_mask = l_idx != np.searchsorted(A[:,0],B[:,0],sorter = sidx,side='right')

# Initialize output array with NaNs. 
# Use l_idx to set rows from A into output array. Use valid_mask to select 
# indices from l_idx and output rows that are to be set.
out = np.full((B.shape[0],A.shape[1]),np.nan)
out[valid_mask] = A[sidx[l_idx[valid_mask]]]

请注意 valid_mask 也可以使用 np.in1d 创建:np.in1d(B[:,0],A[:,0] ) 以获得更直观的答案。但是，我们正在使用 np.searchsorted，因为它在性能方面更好，在 this other solution 中也有更详细的讨论。 .

sample 运行-

In [184]: A
Out[184]: 
array([[45, 11, 86],
       [18, 74, 59],
       [30, 68, 13],
       [55, 47, 78]])

In [185]: B
Out[185]: 
array([[45, 11, 88],
       [55, 83, 46],
       [95, 87, 77],
       [30,  9, 37],
       [14, 97, 98],
       [18, 48, 53]])

In [186]: out
Out[186]: 
array([[ 45.,  11.,  86.],
       [ 55.,  47.,  78.],
       [ nan,  nan,  nan],
       [ 30.,  68.,  13.],
       [ nan,  nan,  nan],
       [ 18.,  74.,  59.]])