Как я могу использовать itertools для вложенного цикла?

Question

Я ищу способ использования itertools.accumulate в starmap.

Я попытался вычислить накопленную сумму каждой строки в таблице, а затем объединить результаты в массив:

# my input
my_table = [[3, 5], [1, 4, 7], [2]]
# expected output
P = [3, 8, 1, 5, 12, 2]

Я использую itertools в цикле for, но он становится намного медленнее, чем другие способы. Так можно ли использовать starmap или другой метод itertools, чтобы сделать это быстрее?

def getSums(my_table):
    P = []
    for i in range(len(my_table)):
       P.append(itertools.accumulate(my_table[i]))
    P = itertools.chain.from_iterable(P)
    return P

Answer 1

Вы можете использовать свой первоначальный циклический подход в значительно упрощенной версии, используя extend, который может принимать любые итерации:

P = []
for lst in my_table:
    P.extend(accumulate(lst))

Answer 2

Использование Itertools хорошо работает, вот что происходит под капотом, можно написать ваши собственные решения просто используют генераторы (без stdlib).

def chain_cumsum(table):
    for it in table:
        yield from cumsum(it)

# write a function to get accumulated sum
def cumsum(l):
    total = 0
    for i in l:
        total += i
        yield total

# then you can get your output:
list(chain_cumsum(my_table))

Answer 3

Вам не нужно starmap, просто используйте встроенную функцию map и скомпонуйте результат, используя itertools.chain():

In [47]: list(chain.from_iterable(map(accumulate, my_table)))
Out[47]: [3, 8, 1, 5, 12, 2]