How can I use itertools for a nested loop?

Question

I'm searching for a method to use itertools.accumulate in starmap.

I tried to calculate the accumulated sum of each row in a table, then concatenate the results to an array:

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

I use itertools in a for loop, but it becomes much slower than other ways. So is it possible to use starmap or other itertools method make it quicker?

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

Kasravnd · Accepted Answer · 2019-01-11 12:49:13Z

3

You don't need starmap just use built-in map function and chain the result using itertools.chain():

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

answered Jan 11, 2019 at 12:49

Kasravnd

108k19 gold badges167 silver badges195 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

leakey Over a year ago

Thanks a lot! I've tried your code and it works well.

user2390182 · Accepted Answer · 2019-01-11 13:17:32Z

1

You can use your initial loop approach in a much simplified version, using extend which can take any iterable:

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

answered Jan 11, 2019 at 13:17

user2390182

73.7k6 gold badges71 silver badges95 bronze badges

Comments

minglyu · Accepted Answer · 2019-01-11 13:25:31Z

1

Use Itertools works well, here is what happens under the hood, you can write your own solutions just use generators(no stdlib needed).

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))

edited Jan 11, 2019 at 13:25

answered Jan 11, 2019 at 13:19

minglyu

3,3872 gold badges18 silver badges38 bronze badges

Collectives™ on Stack Overflow

How can I use itertools for a nested loop?

3 Answers 3

1 Comment

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related