根据二进制数组计算区间之和

Question

我有两个矩阵：

1. 二进制 A = [[1, 0, 1, 0], [0, 0, 1, 0]];
2. 值矩阵 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 的帮助下）？

谢谢

Answer 1

不熟悉python，但我认为你不需要那么多行

define halfSum(matrixA, matrixB):
    sum = 0;
    for i in range(len(matrixA)):
        if matrixA[i] == 1:
           sum += matrixB[i]
        return sum;

Answer 2

你可以使用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])]