Skip to content
Advertisement

Fastest way to split a list into a list of lists based on another list of lists

Say I have a list that contains 5 unique integers in the range of 0 to 9.

import random
lst = random.sample(range(10), 5)

I also have a list of lists, which is obtained by splitting integers from 0 to 19 into 6 groups:

partitions = [[8, 12], [2, 4, 16, 19], [1, 6, 7, 13, 14, 17], [3, 15, 18], [5, 9, 10, 11], [0]]

Now I want to split lst based on the reference partitions. For example, if I have

lst = [0, 1, 6, 8, 9]

I expect the output to be a list of lists like this:

res = [[0], [1, 6], [8], [9]]

I want the algorithm to be as fast as possible. Any suggestions?

Advertisement

Answer

res=[]

for sublist in partitions: # go through all sublists in partitions
    match = [i for i in lst if i in sublist] # find matching numbers in sublist and lst
    if match: # if it is empty don't append it to res
        res.append(match)
# at this point res is [[8], [1, 6], [9], [0]]                                                                                                         
print(sorted(res)) # use sorted to get desired output
User contributions licensed under: CC BY-SA
9 People found this is helpful
Advertisement