2
\$\begingroup\$

I have a list of lists of numbers. The numbers are in sequential order, ranging from 0 to n and each number only appears once. I want to loop over all the numbers in the list of lists and return the index of the list containing that number.

list_of_lists = [[2,3,4,5,6],[0,1,7]]

desired_output = [1,1,0,0,0,0,0,1]

I came up with something, but I'm thinking there has to be a way to do this without doing a list comprehension for each number. Performance is an issue

def get_indices(list_of_lists):

    indices = [0]*len([n for l in list_of_lists for n in l])

    for i in range(len(indices)):
        indices[i] = [ k for k,n in enumerate(list_of_lists) if i in n][0]

    return indices
\$\endgroup\$
2
\$\begingroup\$

The most straightforward approach would be

    for listno, alist in enumerate(list_of_lists):
        for n in alist:
            indices[n] = listno

It surely gives better performance. Is it enough, I don't know.

\$\endgroup\$
1
\$\begingroup\$

It's also a bit faster and easier to read to just get a sum of all the list's lengths instead of the way you're doing it. Building a flattened list just to get its sum is confusing at best and inefficient at worst. Getting a sum of the lists' lengths is much clearer on what your intent is and made the script run faster for me.

indices = [0] * sum(len(each_list) for each_list in list_of_lists)

Or even better, use:

indices = [0] * sum(map(len, list_of_lists))

As Gareth Rees suggested. map will basically run a function on each element in a list, resulting in a list of their lengths which can then be parsed by sum.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.