根据二进制数组

时间:2018-03-01 08:07:26

标签: python-3.x algorithm numpy

我有两个矩阵:

  1. 二进制A = [[1,0,1,0],[0,0,1,0]];
  2. 值矩阵B = [[100,200,300,400],[400,300,100,200]];
  3. 我想计算由矩阵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的帮助下)?

    由于

2 个答案:

答案 0 :(得分:0)

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

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

答案 1 :(得分:0)

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