Compare two lists A, B in Python find all elements in A which correspond to the same number in B

Question

I want to compare two Python lists, \'A\' and \'B\' in such a manner that I can find all elements in A which correspond to the same number in B

Answer 1

Try this approach.

unique = set(B)    # Creates a set of unique entries in B
for u in unique:
    # Find indices of unique entry u
    indices = [i for i, x in enumerate(B) if x == u]

    # Pull out these indices in A
    corrEntry = [A[i] for i in indices]  

    # Do something with the data, in this case print what OP wants
    print('{} corresponds to the number {} of list B'.format(corrEntry , B[indices[0]]))

It finds the unique entries in B by using the set function. We then loop through these unique entries. The first list comprehension (for indices) finds the indices of the entries in B that match this unique entry. The second saves the value in A of those indices.

Answer 2

A = [5, 7, 9, 12, 8, 16, 25]
B = [2, 1, 3, 2, 3, 1, 4]

Create a specific function that takes two lists (A, B) and a number (n) as arguments. Select all items in A that have the same list position as the items in B that are equivalent to n. zip is used to pair items from A and B with the same list position. This function uses a list comprehension to select the items from A.

>>> def f(A, B, n):
    return [a for a, b in zip(A,B) if b == n]

>>> f(A, B, 2)
[5, 12]
>>> 

The function could be written without a list comprehension:

def g(A, B, n):
    result = []
    for a, b in zip(A, B):
        if b == n:
            result.append(a)
    return result

Using fuctools.partial the list arguments can be fixed:

import functools
f_AB = functools.partial(f, A, B)

Then it could be used like this:

>>> f_AB(3)
[9, 8]

>>> numbers = [3, 4, 2]
>>> for n in numbers:
    print (n, f_AB(n))

(3, [9, 8])
(4, [25])
(2, [5, 12])
>>>

Answer 3

An alternative using collections.defaultdict:

import collections as ct

dd = ct.defaultdict(list)
for a, b in zip(A, B):
    dd[b].append(a)

dd
# defaultdict(list, {1: [7, 16], 2: [5, 12], 3: [9, 8], 4: [25]})

Sample of printed results:

for k, v in sorted(dd.items()):
    print("{} corresponding to the number {} of listB".format(v, k))

# [7, 16] corresponding to the number 1 of listB
# [5, 12] corresponding to the number 2 of listB
# [9, 8] corresponding to the number 3 of listB
# [25] corresponding to the number 4 of listB

Answer 4

You can use zip to create tuples which consist from one element from both lists, then sort them and finally group them by value from B:

>>> from itertools import groupby
>>> A = [5, 7, 9, 12, 8, 16, 25]
>>> B = [2, 1, 3, 2, 3, 1, 4]
>>> for k, g in groupby(sorted(zip(B,A)), key=lambda x: x[0]):
...     print('{} corresponds to {}'.format([x[1] for x in g], k))
... 
[7, 16] corresponds to 1
[5, 12] corresponds to 2
[8, 9] corresponds to 3
[25] corresponds to 4

In above zip(B, A) returns iterable of tuples where each tuple has element from B and A:

>>> list(zip(B,A))
[(2, 5), (1, 7), (3, 9), (2, 12), (3, 8), (1, 16), (4, 25)]

Result of above is then sorted so that all the tuples with same value from B are next to each other:

>>> sorted(zip(B,A))
[(1, 7), (1, 16), (2, 5), (2, 12), (3, 8), (3, 9), (4, 25)]

Result of sorting is passed to groupby which groups the tuples based on value returned by key function, in this case the first item in the tuple. Result is iterable of (key, group) tuples where group is iterable of elements:

>>> [(k, list(g)) for k, g in groupby(sorted(zip(B,A)), key=lambda x: x[0])]
[(1, [(1, 7), (1, 16)]), (2, [(2, 5), (2, 12)]), (3, [(3, 8), (3, 9)]), (4, [(4, 25)])]

Answer 5

I think you can use this code:

A = [5, 7, 9, 12, 8, 16, 25]
B = [2, 1, 3, 2, 3, 1, 4]

d = {}

for a, b in zip(A, B):
    d.setdefault(b, [])
    d[b].append(a)

for k, v in sorted(d.items()):
    print('{} corresponds {}'.format(v, k))

Each key of the dictionary will be an element of B, and its associated value will be the list you want.