CODEWITHSUNDEEP
pythonlistmatrix
Morteza R
Morteza R

Reputation: 2319

matrix manipulation using list comprehension

My assignment was to write a Python script to add equal-sized multidimensional matrices together and return the result. For example [[3, 6, 7], [2, 7, -4], [0, 1, 0]] and [[1, 0, 0], [-2, 1, 0], [-1, 1, 3]] should return [[4, 6, 7], [0, 8, -4], [-1, 2, 3]]

I wrote it this way and it works fine:

def addition(a, b):
    if len(a) == len(b):
        c = [[] for i in range(len(a))]
        for i in range(len(a)):
            for j in range(len(a[i])):
                c[i].append(a[i][j] + b[i][j])
        return c
    else:
        print('error')

Next I wondered if it is possible to rewrite this more succinctly using list comprehension and here's my first attempt:

def addition(a, b):
    return [[a[i][j] + b[i][j]] for i in range(len(a)) for j in range(len(a[i]))]

however, with the same values for a and b I get the following result:

>>> a = [[3, 6, 7], [2, 7, -4], [0, 1, 0]]
>>> b = [[1, 0, 0], [-2, 1, 0], [-1, 1, 3]]
>>> addition(a, b)
[[4], [6], [7], [0], [8], [-4], [-1], [2], [3]]

Upvotes: 0

Views: 1141

Answers (2)

khelwood
khelwood

Reputation: 59113

If you move your brackets and indices around a little, you can get what you want:

def addition(a,b):
    return [[a[i][j]+b[i][j] for j in range(len(a[i]))] for i in range(len(a))]

(Or you can use zip, as Cyber suggests.)

Upvotes: 3

Cory Kramer
Cory Kramer

Reputation: 117876

When you are going to do element-wise operations like this, you can use zip to operate sequentially through the lists.

def addition(a,b):
    return [[i+j for i,j in zip(sub1, sub2)] for sub1, sub2 in zip(a, b)]

>>> l1 = [[3, 6, 7], [2, 7, -4], [0, 1, 0]]
>>> l2 = [[1, 0, 0], [-2, 1, 0], [-1, 1, 3]]

>>> addition(l1, l2)
[[4, 6, 7], [0, 8, -4], [-1, 2, 3]]

Upvotes: 4

Related Questions