Search code examples
pythonlistsetindices

how to look for an element in nested list in parallel


I have a list of elements A = ['a', 'b', 'c', 'd'] and a nested list B = [{'55','a'}, {'4','t'}, {'x','y','zx'}, {'b','d'}, {'66','c'}].
In addition, each element can appear only in one sub-list of B.

I want to find, for each element from A, the index of the sub-list from B that contains it ([0, 3, 4, 3]). Assume that all the elements from A appear somewhere in B.

Using a nested loop isn't practical for the size of my data.


Solution

  • First build a dictionary of the positions in B for each item, then use a list comprehension on A:

    d = {k:i for i,s in enumerate(B) for k in s}
    # {'55': 0, 'a': 0, 't': 1, '4': 1,
    #  'zx': 2, 'y': 2, 'x': 2, 'd': 3,
    #  'b': 3, 'c': 4, '66': 4}
    
    out = [d[x] for x in A]
    

    Or, if you are not sure that all elements of A are in B:

    out = [d.get(x, -1) for x in A]
    

    Output:

    [0, 3, 4, 3]
    

    Used input:

    A = ['a', 'b', 'c', 'd']
    B = [{'55','a'}, {'4','t'}, {'x','y','zx'}, {'b','d'}, {'66','c'}]