首页 > 解决方案 > 根据条件加入两个不同大小的数组

问题描述

我有两个不同大小的数组,我希望它们只有在前三个元素匹配时才合并。

例如:

A1 = [
    [1,8,9,8,2],
    [2,9,9,8,2],
    [1,7,7,8,2],
    [8,6,2,6,7]
]


A2 = [[2,9,9,30,60],
      [8,6,2,70,20]]

Result = [[2,9,9,8,2,30,60]
         [8,6,2,6,7,70,20]]

我想这有一些 numpy 函数,但我找不到它。

标签: pythonnumpy

解决方案


result = [i+j[3:] for i in A1 for j in A2 if i[:3]==j[:3]]

这相当于:

result = []
for i in A1:
    for j in A2:
       if (i[:3] == j[:3]): # if the first three element match

           # skip the first 3 elements in j (they're already included in i)
           result.append(i + j[3:]) 

numpy用数组做到这一点

result = np.array([np.concatenate((i,j[3:])) for i in A1 for j in A2 if all(i[:3]==j[:3])])

这里我们用来np.concatenate()合并两个数组,用来all()比较前3个元素。


推荐阅读