矩阵（列表）的元素总和，但仅在另一个矩阵（列表）中给出的索引处

Question

我有一些包含整数的矩阵：

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

和一个矩阵包含：

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

我想对矩阵的所有元素求和，其中索引中的元素等于-1，但我不能使用任何包。 预期输出应为 935。

提前致谢。

Answer 1

使用生成器表达式：

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

输出

935

作为备选：

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

Answer 2

您可以使用 itertools 将矩阵展平并沿索引压缩它::

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