Numpy add array of shape NxD to NxK into the first D in array NxK-CodePudding

I have two arrays

a = np.array([[0, 0, 2, 3, 4],
              [0, 1, 2, 3, 4],
              [0, 2, 2, 3, 4]])

and

b = np.array([[1, 1],
              [2, 2],
              [3, 3]])

I want to one array where I am adding the values of b to the first two columns in a like this:

c = np.array([[1, 1, 2, 3, 4],
              [2, 3, 2, 3, 4],
              [3, 5, 2, 3, 4]])

if it helps you can think of the first two columns in a as the x,y coordinates and b as dx, dy.

My current method is as follows:

c = np.concatenate([a[:, 0:2]   b, a[:, 2:]],1)

but I am looking for a better method

Thank you

CodePudding user response：

You can use np.pad to add zeros to b to make its shape the same as a's, then add them:

>>> a   np.pad(b, ((0, 0), (0, 3)))
array([[1, 1, 2, 3, 4],
       [2, 3, 2, 3, 4],
       [3, 5, 2, 3, 4]])

In general (for 2-D):

>>> a = np.array([[0, 0, 2, 3, 4],
...               [0, 1, 2, 3, 4],
...               [0, 2, 2, 3, 4]])
>>> b = np.array([[1, 1],
...               [2, 2],
...               [3, 3],
...               [4, 4],
...               [5, 5]])
>>> a_shape, b_shape = a.shape, b.shape
>>> max_w = max(a_shape[0], b_shape[0])
>>> max_h = max(a_shape[1], b_shape[1])
>>> padded_a = np.pad(a,
                      ((0, np.abs(a_shape[0] - max_w)),
                       (0, np.abs(a_shape[1] - max_h))))
>>> padded_a
array([[0, 0, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 2, 2, 3, 4],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]])
>>> padded_b = np.pad(b,
                      ((0, np.abs(b_shape[0] - max_w)),
                       (0, np.abs(b_shape[1] - max_h))))
>>> padded_b
array([[1, 1, 0, 0, 0],
       [2, 2, 0, 0, 0],
       [3, 3, 0, 0, 0],
       [4, 4, 0, 0, 0],
       [5, 5, 0, 0, 0]])
>>> padded_a   padded_b
array([[1, 1, 2, 3, 4],
       [2, 3, 2, 3, 4],
       [3, 5, 2, 3, 4],
       [4, 4, 0, 0, 0],
       [5, 5, 0, 0, 0]])

In general (2-D, using a zeros array and adding to it):

>>> c = np.zeros((max_h, max_w), dtype=a.dtype)
>>> c[:a_shape[0], :a_shape[1]]  = a
>>> c[:b_shape[0], :b_shape[1]]  = b
>>> c
array([[1, 1, 2, 3, 4],
       [2, 3, 2, 3, 4],
       [3, 5, 2, 3, 4],
       [4, 4, 0, 0, 0],
       [5, 5, 0, 0, 0]])