Search code examples
pythonlistmatrixsum

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

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.

Solution

  • Use a generator expression:

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

    Output

    935

    As an alternative:

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