CODEWITHSUNDEEP
pythonnumpy
GoingMyWay
GoingMyWay

Reputation: 17478

Numpy add given rows

Given a numpy array for example

A = np.ones(shape=(7, 6), dtype=np.float32)

and a list 

v = [[0, 2], [1, 4], [3, 5, 6]]

what I want to do is adding rows in A given each item in v, for v, there are 3 items, for v[0], add row 0 and row 2 column-wisely. The shape of the output is (3, 6) and the output is 

res = array([[2., 2., 2., 2., 2., 2.],
             [2., 2., 2., 2., 2., 2.],
             [3., 3., 3., 3., 3., 3.]])

# res[0] = A[0] + A[2]
# res[1] = A[1] + A[4]
# res[2] = A[3] + A[5] + A[6]

Here is a more clear example, give a matrix

m = [[1, 2, 3],
      [2, 3, 4],
      [1, 1, 1],
      [2, 2, 2],
      [1, 1, 1]]

and rows to be added

 v = [[0, 1, 3], [2]]

so, here add rows 0, 1, and 3 in matrix m and since there is only one row to be added in [2], so the result is 

# res.shape = (2, 3)
res[0] = m[0] + m[1] + m[3]
res[1] = m[2]

Are there any more elegant way to do so?

Upvotes: 2

Views: 40

Answers (1)

timgeb
timgeb

Reputation: 78780

You can use fancy indexing to select rows from your array.

For A:

>>> A = np.ones(shape=(7, 6), dtype=np.float32)
>>> v = [[0, 2], [1, 4], [3, 5, 6]]
>>> np.array([A[rows].sum(axis=0) for rows in v])
array([[2., 2., 2., 2., 2., 2.],
       [2., 2., 2., 2., 2., 2.],
       [3., 3., 3., 3., 3., 3.]], dtype=float32)

For m:

>>> m = np.array([[1, 2, 3], [2, 3, 4], [1, 1, 1], [2, 2, 2], [1, 1, 1]])
>>> v = [[0, 1, 3], [2]]
>>> np.array([m[rows].sum(axis=0) for rows in v])
array([[5, 7, 9],
       [1, 1, 1]])

I don't know if this can be vectorized further.

Upvotes: 2

Related Questions