Python: find the index of sublist that contains an element

Question

I have a list of sublists. I want to find the index of the sublist that contains the desired element.

E.g.

aa = [[1,2,3],[4,5,6],[7,8,9]]

When I search for element 7, it should return 2.

How can I find the index?

Answer 1

There are multiple ways to find the index of the sublist that contains the desired element.

Approach 1:

You can iterate over sublists to check if it contains the desired element. This will return a list of True and False. Then you can find the index of True which will be the index of the sublist.

Here is an example:

>>> aa = [[1,2,3],[4,5,6],[7,8,9]]
>>> n=7
>>> [n in a for a in aa]
[False, False, True]
>>> [n in a for a in aa].index(True)
2

Approach 2:

You can also the isin() function of NumPy.  The NumPy will treat the list of sublists as a 2D array and the function will return a 2D array of True and False. Using the nonzero() function, you can find the row and column number of True in the array. The row number will be the index of the sublist containing the desired element.

Here is an example:

>>> import numpy as np
>>> aa = [[1,2,3],[4,5,6],[7,8,9]]
>>> t=np.isin(aa,6)
>>> t
array([[False, False, False],
       [False, False,  True],
       [False, False, False]])
>>> np.nonzero(t)
(array([1], dtype=int64), array([2], dtype=int64))
>>> np.nonzero(t)[0][0]
1