Multidimensional list match in python

Question

This has caused some serious headache today. Suppose I have two instances of my object, instance A and instance B. These come with properties is the form of a list. Say the two properties for A are

a1 = [1, 2, 3, 4, 5]
a2 = [10, 20, 30, 40, 50]

and those for B:

b1 = [5, 7, 3, 1]
b2 = [50, 20, 30, 20]

What I want is to simply find the indices in b1 and b2, where a pair equals the values in a1 and a2. So in this example this would be the indices 0 and 2 since for those we have

b1[0] = 5 and b2[0] = 50

which we find in a1 and a2 as the last entries. Same for index 2 for which we find (3, 30) in (b1, b2) which is also in (a1, a2).

Note here, that the lists a1 and a2 have always the same length as well as b1 and b2.

Any help? 😊

Answer 1

You can use a combination of zip , set and enumerate :

>>> a1 = [1, 2, 3, 4, 5]
>>> a2 = [10, 20, 30, 40, 50]
>>> b1 = [5, 7, 3, 1]
>>> b2 = [50, 20, 30, 20]
>>> a12 = set(zip(a1, a2))
>>> [i for i, e in enumerate(zip(b1, b2)) if e in a12]
[0, 2]

With zip , you group the pairs together, and with set you turn them into a set, as order does not matter and set have faster lookup. Then, enumerate gives you pairs of indices and elements, and using the list-comprehension you get those indices from b12 whose elements are in a12 .

Answer 2

I think another structure would be better?

a tuple, or a key set ...

a = [(1,10),(2,20)] and so on

edit

well... tobias_k shows you how :)

Answer 3

Try this

In [38]: [b1.index(i[0]) for i in zip(a1,a2) for j in zip(b1,b2) if i==j]
Out[38]: [2, 0]

Answer 4

There is also the possibility to check for each element in (a1, a2) whether it is in (b1, b2) and it will return all matches in a list and will take care of duplicates:

a1 = [1, 2, 3, 4, 5]
a2 = [10, 20, 30, 40, 50]

b1 = [5, 7, 3, 1, 5]
b2 = [50, 20, 30, 20, 50]

# Construct list of tuples for easier matching
pair_a = [(i, k) for i, k in zip(a1, a2)]
pair_b = [(i, k) for i, k in zip(b1, b2)]

# Get matching indices (for each entry in pair_a get the indices in pair_b)
indices = [[i for i, j in enumerate(pair_b) if j == k] for k in pair_a]

gives

[[], [], [2], [], [0, 4]]