我有两个矩阵:
- 二进制A = [[1,0,1,0],[0,0,1,0]];
- 值的矩阵b = [[100,200,300,400],[400,300,100,200]];
我想计算由矩阵A的行形成的间隔的总和。结果将如下:r = [[[300,0,700,0],[0,0,300,0]](通常,不必设置零[[300,700],[300]] - 这也是正确的解决方案(
我已经编写了代码,但是非常可怕(尽管工作正常(
def find_halfsum(row1, row2):
i = 0
result = []
count = 0
for j in range(len(row1)):
if row1[j] == 1 and count == 0:
i = j
count += 1
elif row1[j] == 1:
count += 1
if count == 2:
if j == i + 1:
result.append(row2[i])
else:
result.append(sum(row2[i:j]))
i = j
count = 1
if j == len(row1) - 1:
result.append(sum(row2[i:j + 1]))
return result
有人知道美丽的解决方案(最好更快((最好借助numpy(?
谢谢
不熟悉python,但我认为您不需要那么多行
define halfSum(matrixA, matrixB):
sum = 0;
for i in range(len(matrixA)):
if matrixA[i] == 1:
sum += matrixB[i]
return sum;
您可以使用numpy.add.reduceat
:
>>> A = np.array([[1, 0, 1, 0], [0, 0, 1, 0]])
>>> B = np.array([[100, 200, 300, 400], [400, 300, 100, 200]])
>>>
>>> [np.add.reduceat(b, np.flatnonzero(a)) for a, b in zip(A, B)]
[array([300, 700]), array([300])]