yatu
yatu

Reputation: 88275

Concatenate with broadcast

Consider the following arrays:

a = np.array([0,1])[:,None]
b = np.array([1,2,3])

print(a)
array([[0],
       [1]])

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

Is there a simple way to concatenate these two arrays in a way that the latter is broadcast, in order to obtain the following?

array([[0, 1, 2, 3],
       [1, 1, 2, 3]])

I've seen there is this closed issue with a related question. An alternative is proposed involving np.broadcast_arrays, however I cannot manage to adapt it to my example. Is there some way to do this, excluding the np.tile/np.concatenate solution?

Upvotes: 9

Views: 4638

Answers (3)

Frank Yellin
Frank Yellin

Reputation: 11297

I had a similar problem where I had two matrices x and y of size (X, c1) and (Y, c2) and wanted the result to be the matrix of size (X * Y, c1 + c2) where the rows of the result were all the concatenations of rows from x and rows from y.

I, like the original poster, was disappointed to discover that concatenate() would not do broadcasting for me. I thought of using the solution above, except X and Y could potentially be large, and that solution would use a large temporary array.

I finally came up with the following:

    result = np.empty((x.shape[0], y.shape[0], x.shape[1] + y.shape[1]), dtype=x.dtype)
    result[...,:x.shape[0]] = x[:,None,:]
    result[...,x.shape[0]:] = y[None,:,:]
    result = result.reshape((-1, x.shape[1] + y.shape[1]))

I create a result array of size (X, Y, c1 + c2), I broadcast in the contents of x and y, and then reshape the results to the right size.

Upvotes: 0

Christoph Boeddeker
Christoph Boeddeker

Reputation: 674

Here a more general solution:

def concatenate_broadcast(arrays, axis=-1):
    def broadcast(x, shape):
        shape = [*shape]  # weak copy
        shape[axis] = x.shape[axis]
        return np.broadcast_to(x, shape)

    shapes = [list(a.shape) for a in arrays]
    for s in shapes:
        s[axis] = 1

    broadcast_shape = np.broadcast(*[
            np.broadcast_to(0, s)
            for s in shapes
    ]).shape

    arrays = [broadcast(a, broadcast_shape) for a in arrays]
    return np.concatenate(arrays, axis=axis)

Upvotes: 1

Anubhav Natani
Anubhav Natani

Reputation: 344

You can do it in the following way

import numpy as np
a = np.array([0,1])[:,None]
b = np.array([1,2,3])
b_new = np.broadcast_to(b,(a.shape[0],b.shape[0]))
c = np.concatenate((a,b_new),axis=1)
print(c)

Upvotes: 8

Related Questions