Get list of indices from searching for values in list of lists

Question

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

Answer 1

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.

Answer 2

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.