assigning values in a numpy array assigning values in a numpy array numpy numpy

assigning values in a numpy array


At the moment, I can only think of the “simple” version, which involves flattening along the first two dimensions. This code should work:

shape_last = x.shape[-1]x.reshape((-1, shape_last))[np.arange(y.size), y.flatten()] = 1

This yields (with my randomly-generated y):

array([[[ 0.,  0.,  0.,  1.],        [ 0.,  0.,  1.,  0.],        [ 0.,  1.,  0.,  0.]],       [[ 0.,  1.,  0.,  0.],        [ 0.,  0.,  0.,  1.],        [ 0.,  1.,  0.,  0.]]])

The key is, if you do an indexing using multiple numpy arrays (advanced indexing), numpy will use pairs of indices to index into the array.

Of course, make sure x and y are both either C-order or F-order — otherwise, the calls to reshape and flatten might give different orders.


Use numpy.meshgrid() to make arrays of indexes that you can use to index into both your original array and the array of values for the third dimension.

import numpy as npimport scipy as spimport scipy.stats.distributionsa = np.zeros((2,3,4))z = sp.stats.distributions.randint.rvs(0, 4, size=(2,3))xx, yy = np.meshgrid( np.arange(2), np.arange(3) )a[ xx, yy, z[xx, yy] ] = 1print a

I've renamed your array from x to a, and the array of indexes from y to z, for clarity.

EDIT: 4D example:

a = np.zeros((2,3,4,5))z = sp.stats.distributions.randint.rvs(0, 4, size=(2,3))w = sp.stats.distributions.randint.rvs(0, 5, size=(2,3))xx, yy = np.meshgrid( np.arange(2), np.arange(3) )a[ xx, yy, z[xx, yy], w[xx, yy] ] = 1


x = np.zeros((2,3,4))y=np.array([[2, 1, 0],[3, 2, 0]]) # or y=sp.stats...for i in range(2):    for j in range(3):        x[i,j,y[i,j]]=1

will produce the desired result, IIRC. If the array dimensions never change, consider replacing the two for loops and their burden by

for j in range(3):    x[0,j,y[0,j]] = x[1,j,y[1,j]] = 1