How to fill values of a (2,3,4) array using a (2,3) array of indices and a (2,3) array of new values

Viewed 36

I have the following arrays.

b0 = np.zeros((2, 3, 4))
array([[[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],
       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]]])

b1 = np.array([
    [1, 0, 2],
    [0, 0, 1]
])

b2 = np.array([
    [100, 0, 0],
    [200, 300, 0]
])

My goal is to update values of b0 using indices specified in b1 and new values given by b2. Specifically, I want to make these updates

b0[0, 0, b1[0,0]] = b2[0,0]
b0[0, 1, b1[0,1]] = b2[0,1]
b0[0, 2, b1[0,2]] = b2[0,2]

b0[1, 0, b1[1,0]] = b2[1,0]
b0[1, 1, b1[1,1]] = b2[1,1]
b0[1, 2, b1[1,2]] = b2[1,2]

b0
array([[[  0, 100,   0,   0],
        [  0,   0,   0,   0],
        [  0,   0,   0,   0]],
       [[200,   0,   0,   0],
        [300,   0,   0,   0],
        [  0,   0,   0,   0]]])

How do I do this dynamically?

3 Answers

Just realized I can do this

n = b0.shape[0]
k = b2.shape[1]
i1 = np.repeat(np.arange(n), k)
i2 = np.tile(np.arange(k), n)
b0[i1, i2, b1[i1, i2]] = b2[i1, i2]

but there's probably a better way to do this.

In [51]: b1 = np.array([
    ...:     [1, 0, 2],
    ...:     [0, 0, 1]
    ...: ])
    ...: 
    ...: b2 = np.array([
    ...:     [100, 0, 0],
    ...:     [200, 300, 0]
    ...: ])
In [52]: b0 = np.zeros((2,3,4),int)
In [53]: b0[np.arange(2)[:,None], np.arange(3), b1] = b2
In [54]: b0
Out[54]: 
array([[[  0, 100,   0,   0],
        [  0,   0,   0,   0],
        [  0,   0,   0,   0]],

       [[200,   0,   0,   0],
        [300,   0,   0,   0],
        [  0,   0,   0,   0]]])

np.arange(2)[:,None], np.arange(3) broadcast together to index a (2,3) block, same as b1 shape.

meshgrid lets us compare my indexing and yours:

In [58]: np.meshgrid(np.arange(2), np.arange(3), sparse=True, indexing='ij')
Out[58]: 
[array([[0],
        [1]]),
 array([[0, 1, 2]])]
In [59]: np.meshgrid(np.arange(2), np.arange(3), sparse=False, indexing='ij')
Out[59]: 
[array([[0, 0, 0],
        [1, 1, 1]]),
 array([[0, 1, 2],
        [0, 1, 2]])]

Try:

ind = np.argwhere(~np.isnan(b1))
ind1, ind2 = zip(*ind)

b0[ind1, ind2, b1[ind1,ind2]]=np.ravel(b2)

However I see it's 0 on (0,0,0), while you have 200 there - which I feel might be your overlook (correct me, if I'm wrong though)...

Related