Copying sublist structure to another list of equal length in Python

Question

Say I have one list,

list1 = ['Dog', 'Cat', 'Monkey', 'Parakeet', 'Zebra']

And another,

list2 = [[True, True], [False], [True], [False]]

(You can imagine that the second list was created with an itertools.groupby on an animal being a house pet.)

Now say I want to give the first list the same sublist structure as the second.

list3 = [['Dog', 'Cat'], ['Monkey'], ['Parakeet'], ['Zebra']]

I might do something like this:

list3 = []
lengths = [len(x) for x in list2]
count = 0
for leng in lengths:
    templist = []
    for i in range(leng):   
        templist.append(list1[count])
        count += 1
    list3.append(templist)

Which I haven't totally debugged but I think should work. My question is if there is a more pythonic or simple way to do this? This seems kind of convoluted and I have to imagine there is a more graceful way to do it (maybe in itertools which never ceases to impress me).

Answer 1

You can simply use next for a no-import solution:

list1 = ['Dog', 'Cat', 'Monkey', 'Parakeet', 'Zebra']
list_1 = iter(list1)
list2 = [[True, True], [False], [True], [False]]
new_list2 = [[next(list_1) for _ in i] for i in list2]

Output:

[['Dog', 'Cat'], ['Monkey'], ['Parakeet'], ['Zebra']]

However, for input n levels deep, you can use recursion:

list1 = ['Dog', 'Cat', 'Monkey', 'Parakeet', 'Zebra']
list_1 = iter(list1)
list2 = [[True, [True, False]], [False], [[True]]]
def place_data(current):
   return [next(list_1) if isinstance(i, int) else place_data(i) for i in current]

print(place_data(list2))

Output:

[['Dog', ['Cat', 'Monkey']], ['Parakeet'], [['Zebra']]]

Answer 2

Make an iter of list1 and islice the necessary N elements of each list2,e g:

from itertools import islice

list1 = ['Dog', 'Cat', 'Monkey', 'Parakeet', 'Zebra']
list2 = [[True, True], [False], [True], [False]]

it = iter(list1)
list3 = [list(islice(it, len(el))) for el in list2]
# [['Dog', 'Cat'], ['Monkey'], ['Parakeet'], ['Zebra']]

Answer 3

itertools.groupby allows you to do this already, because it returns the groups as well as the keys

def is_pet(animal):
    return animal.lower() in ('dog', 'cat', 'parakeet')

animals = ['Dog', 'Cat', 'Monkey', 'Parakeet', 'Zebra']
print([list(g) for k, g in itertools.groupby(animals, key=is_pet)])
# [['Dog', 'Cat'], ['Monkey'], ['Parakeet'], ['Zebra']]

The other option that springs to mind is to use an iterator over the list, something like

def group(to_group, group_like):
    it = iter(to_group)
    for sub_list in group_like:
        yield [next(it) for _ in sub_list]

print(list(group(list1, list2)))
# [['Dog', 'Cat'], ['Monkey'], ['Parakeet'], ['Zebra']]

Answer 4

This is one way:

from itertools import accumulate

idx_acc = [0] + list(accumulate(map(len, list2)))
list1 = [list1[i:j] for i, j in zip(idx_acc, idx_acc[1:])]

# [['Dog', 'Cat'], ['Monkey'], ['Parakeet'], ['Zebra']]