pythonlistfor-loop

How do I efficiently and briefly add columns of a partial table?


I want to use native python and write a terse code as much as possible to achieve this. So with table A as below:

| 1 | 2 | 3 |

| 4 | 5 | 6 |

Then the list A is given as following

A = [[1, 2, 3], [4, 5, 6]]

And suppose the task is to sum first two columns of A into a list B.

I tried naive approach of using a nested for loop. (which works)

n = 2
B = [0] * n
for row in range(len(A)):
    for col in range(n):
        B[col] += A[row][col]

I tried list comprehension

B = [sum(x) for x in zip(*A)][:n]

But the problem is that I don't want to go through all the columns and then later split list B. I'd like to avoid computing elements that's going to be later thrown away.

What would be the most efficient and concise way to solve this task?


Solution

  • To avoid computing unnecessary values in the list comprehension, you would want to slice zip(*A), not the final list.

    Since zip returns an iterator, not a list, you can't use [:n].

    But you can use itertools.islice:

    from itertools import islice
    
    B = [sum(x) for x in islice(zip(*A), n)]