Python在特定列上合并两个数组

Question

I have two sets/arrays/lists of the form

a = [(12, 14, 0.3, 0.6, 0.8), (16, 18, 0.4, 0.5, 0.3), (19, 22, 0.4, 0.5, 0.3)]
b = [(12, 14, 44, 12), (5, 4, 66, 12), (19, 22, 96, 45)]

And I would like to find c which is a list of items from b found in a such that only the first two elements in the tuple need to match (ex 12 14). So in this case the answer c would be

c = [(12, 14, 44, 12), (19, 22, 96, 45)]

I used nested loops however it is way too slow. Thanks

Answer 1

You can do this with a list comprehension

>>> a = [(12, 14, 0.3, 0.6, 0.8), (16, 18, 0.4, 0.5, 0.3), (19, 22, 0.4, 0.5, 0.3)]
>>> b = [(12, 14, 44, 12), (5, 4, 66, 12), (19, 22, 96, 45)]
>>> [item for item in b for checker in a if item[:2] == checker[:2]]
[(12, 14, 44, 12), (19, 22, 96, 45)]

Answer 2

You can do this O(N) time if you store all the unique two item tuples from a in a set first:

>>> keys = {x[:2] for x in a}
>>> [x for x in b if x[:2] in keys]
[(12, 14, 44, 12), (19, 22, 96, 45)]

Note that if you're only trying to match items on the same index, then simply use zip with a list comprehension:

>>> [y for x, y in zip(a, b) if x[:2] == y[:2]]
[(12, 14, 44, 12), (19, 22, 96, 45)]

#Equivalent Numpy version:
>>> arr_a = np.array(a)
>>> arr_b = np.array(b)
>>> arr_b[(arr_b[:,:2] == arr_a[:,:2]).all(axis=1)]
array([[12, 14, 44, 12],
       [19, 22, 96, 45]])

Answer 3

you can do with numpy if you're using numpy

In [49]: a = np.array([(12, 14, 0.3, 0.6, 0.8), (16, 18, 0.4, 0.5, 0.3), (19, 22, 0.4, 0.5, 0.3)])

In [50]: b = np.array([(12, 14, 44.0, 12.0), (5, 4, 66.0, 12.0), (19, 22, 96.0, 45.0)])

In [51]: print b[np.all(a[:,:2]==b[:,:2],1)]
[[ 12.  14.  44.  12.]
 [ 19.  22.  96.  45.]]

How it works?

In [52]: print a[:,:2]==b[:,:2]
[[ True  True]
 [False False]
 [ True  True]]

np.all takes an array of booleans and reduces using a logical and along the axis specified by the optional second argument (or using all the elements)

In [53]: print np.all(a[:,:2]==b[:,:2])
False

In [69]: print np.all(a[:,:2]==b[:,:2],1)
[ True False  True]

In [70]: print np.all(a[:,:2]==b[:,:2],0)
[False False]

In [71]:

in our case, of course, the right axis to use is 1 .

(ps: I must confess a bit of sloppiness in treating the types of your array values)

Answer 4

List Comprehension

      >>> a = [(12, 14, 0.3, 0.6, 0.8), (16, 18, 0.4, 0.5, 0.3), (19, 22, 0.4, 0.5, 0.3)]
      >>> b = [(12, 14, 44, 12), (5, 4, 66, 12), (19, 22, 96, 45)]
      >>> c=[j  for i in a for j in b  if i[:2]==j[:2]]

Output:

 [(12, 14, 44, 12), (19, 22, 96, 45)]