通用情况
方法 #1:您可以使用np.add.at
执行这样的ID-based
加法操作,如下所示 -
out_id = np.union1d(a[:,0],b[:,0])
out_count = np.zeros_like(out_id)
_,a_idx = np.where(a[:,None,0]==out_id)
_,b_idx = np.where(b[:,None,0]==out_id)
out_count[a_idx] = a[:,1]
np.add.at(out_count, b_idx,b[:,1])
out = np.column_stack((out_id,out_count))
要查找a_idx
和b_idx
,可能更快的替代方法是使用np.searchsorted
,代码如下 -
a_idx = np.searchsorted(out_id, a[:,0], side='left')
b_idx = np.searchsorted(out_id, b[:,0], side='left')
样例输入输出:
In [538]: a
Out[538]:
array([[1, 2],
[4, 2],
[3, 1],
[5, 5]])
In [539]: b
Out[539]:
array([[3, 7],
[1, 1],
[4, 0],
[2, 3],
[6, 2]])
In [540]: out
Out[540]:
array([[1, 3],
[2, 3],
[3, 8],
[4, 2],
[5, 5],
[6, 2]])
方法 #2:你可以使用np.bincount
执行基于ID的累加操作-
out_id = np.union1d(a[:,0],b[:,0])
id_arr = np.concatenate((a[:,0],b[:,0]))
count_arr = np.concatenate((a[:,1],b[:,1]))
summed_vals = np.bincount(id_arr,count_arr)
mask = np.in1d(np.arange(np.max(out_id)+1),out_id)
out_count = summed_vals[mask]
out = np.column_stack((out_id,out_count))
具体情况
如果a
和b
中的ID列已经排序,那么就更容易了,因为我们可以使用带有np.in1d
的掩码来索引到使用np.union
创建的输出ID数组,如下所示 -
out_id = np.union1d(a[:,0],b[:,0])
mask1 = np.in1d(out_id,a[:,0])
mask2 = np.in1d(out_id,b[:,0])
out_count = np.zeros_like(out_id)
out_count[mask1] = a[:,1]
np.add.at(out_count, np.where(mask2)[0],b[:,1])
out = np.column_stack((out_id,out_count))
样例运行 -
In [552]: a
Out[552]:
array([[1, 2],
[2, 2],
[3, 1],
[4, 5],
[8, 5]])
In [553]: b
Out[553]:
array([[2, 2],
[3, 1],
[4, 0],
[5, 3],
[6, 2],
[8, 2]])
In [554]: out
Out[554]:
array([[1, 2],
[2, 4],
[3, 2],
[4, 5],
[5, 3],
[6, 2],
[8, 7]])