numpy 2D array assignment with 2D value and indices arrays

Question

My goal is to assign the values of an existing 2D array, or create a new array, using two 2D arrays of the same shape, one with values and one with indices to assign the corresponding value to.

X = np.array([range(5),range(5)])
X
array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]])

Y= np.array([range(5), [2,3,4,1,0]])
Y
array([[0, 1, 2, 3, 4],
       [2, 3, 4, 1, 0]])

My desired output is an array of the same shape as X and Y, with the values of X given in the index from the corresponding row in Y. This result can be achieved by looping through each row in the following way:

output = np.zeros(X.shape)

for i in range(X.shape[0]):
    output[i][Y[i]] = X[i]
output 
array([[ 0.,  1.,  2.,  3.,  4.],
       [ 4.,  3.,  0.,  1.,  2.]])

Is there a more efficient way to apply this sort of assignment?

np.take(output, Y)

Will return the items in the output array I would like to assign to the values of X to, but I believe np.take does not produce a reference to the original array, and instead a new array.

This can be done with the sparse matrix definition in scipy. Have a look at coo_matrix definition in docs.scipy.org/doc/scipy-0.15.1/reference/generated/… . — Bort
– Bort, Commented Apr 10, 2015 at 13:54

unutbu · Accepted Answer · 2015-04-10 14:05:33Z

for i in range(X.shape[0]):
    output[i][Y[i]] = X[i]

is equivalent to

I = np.arange(X.shape[0])[:, np.newaxis]
output[I, Y] = X

For example,

X = np.array([range(5),range(5)])
Y = np.array([range(5), [2,3,4,1,0]])
output = np.zeros(X.shape)
I = np.arange(X.shape[0])[:, np.newaxis]
output[I, Y] = X

yields

>>> output
array([[ 0.,  1.,  2.,  3.,  4.],
       [ 4.,  3.,  0.,  1.,  2.]])

There is not much difference in performance when the loop has few iterations. But if X.shape[0] is large, then using indexing is much faster:

def using_loop(X, Y):
    output = np.zeros(X.shape)
    for i in range(X.shape[0]):
        output[i][Y[i]] = X[i]
    return output

def using_indexing(X, Y):
    output = np.zeros(X.shape)
    I = np.arange(X.shape[0])[:, np.newaxis]
    output[I, Y] = X
    return output

X2 = np.tile(X, (100,1))
Y2 = np.tile(Y, (100,1))

In [77]: %timeit using_loop(X2, Y2)
1000 loops, best of 3: 376 µs per loop

In [78]: %timeit using_indexing(X2, Y2)
100000 loops, best of 3: 15.2 µs per loop

Collectives™ on Stack Overflow

numpy 2D array assignment with 2D value and indices arrays

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related