Sum of elements of a matrix (list), but only at indexes given in another matrix (list)-CodePudding

I have some matrix containing integers:

matrix = [[85, 61, 48, 100, 96],
 [23, 72, 13, 45, 36],
 [97, 80, 65, 84, 46],
 [80, 59, 76, 61, 99]]

and a matrix containing:

index = [[-1, 0, 0, -1, -1],
 [0, -1, 0, -1, -1],
 [-1, 0, -1, -1, 0],
 [-1, 0, -1, 0, -1]]

I want to sum all the elements of matrix, where the element in index is equal to -1, but I can't use any package. The expected output should be 935.

Thanks in advance.

CodePudding user response：

res = sum(v for ii, val in zip(index, matrix) for i, v in zip(ii, val) if i == -1)
print(res)

Output

As an alternative:

res = 0
for ii, val in zip(index, matrix):
    for i, v in zip(ii, val):
        if i == -1:
            res  = v

CodePudding user response：

You could use the itertools to flatten the matrix and zip it along the indices::

import itertools as it
sum([v for (i, v) in zip(it.chain.from_iterable(index), it.chain.from_iterable(matrix)) if i == -1])