Применить операцию и список к вложенному списку

Question

Существует большое количество вопросов об итерациях по вложенным спискам, однако я бы хотел перебрать вложенный список и применить другой список.

Вот мой сценарий:

def operation(a, b):
    return a + b

def magic_function(func, list, nested_list):
    """iterate over nested_list and apply the next element of list"""
    ...

magic_function(operation, [0, 0, 0, 10, 10, 10], [[1, 2, 3], [1, 2, 3]])

---

Желаемый вывод:

[[1, 2, 3], [11, 12, 13]]

Может быть сильное желание ответить на это недоуменным, но в реальном сценарии это объекты, а не числа.

Стандарт itertools.chain.from_iterable здесь не работает, так как он не сохраняет вложенность списка.

Answer 1

Вот реализация, которая работает с произвольно глубоко вложенными списками:

def nested_map(f, obj):
    if isinstance(obj, list):
        return [nested_map(f, item) for item in obj]
    else:
        return f(obj)

def magic_function(func, seq, nested_list):
    g = iter(seq)
    return nested_map(lambda item: func(item, next(g)), nested_list)

def operation(a, b):
    return a + b

result = magic_function(operation, [0, 0, 0, 10, 10, 10], [[1, 2, 3], [1, 2, [[3]]]])
print(result)

Результат:

[[1, 2, 3], [11, 12, [[13]]]]

Answer 2

Вы можете использовать itertools.count для индексации своего плоского списка:

from itertools import count

def operation(a, b):
    return a + b

def foo(func, lst, nested_list):
    idx = count()
    res = [[operation(j, lst[next(idx)]) for j in i] for i in nested_list]
    return res

res = foo(operation, [0, 0, 0, 10, 10, 10], [[1, 2, 3], [1, 2, 3]])

print(res)

[[1, 2, 3], [11, 12, 13]]

Или вы можете перебрать плоский список напрямую:

def foo(func, lst, nested_list):
    k = iter(lst)
    res = [[operation(j, next(k)) for j in i] for i in nested_list]
    return res