python: broadcast sliced np.array assignment to any number of dimensions

Question

I have np.arrays C, R and S of shapes (?, d), (?, n) and (?, d) respectively; where d<=n and the question mark represents any number of matching dimensions. Now I would like to do the following assignment (this is of course not proper python code, but it works if ? is just a single number):

for i in range(?):
    R[i][S[i]]=C[i]

That is: I want for each tuple i of indices (within the bounds specified by ?) to take the corresponding array R[i] in R and assign d many positions (the ones specified by S[i]) to be the values in the array C[i].

What is the pythonic way to do this?

Example:

setup

import numpy as np
m,n,d= 2,7,4
R=np.zeros((m,n))
C=np.arange(d*m).reshape((m,d))
S=np.array([[0,2,4,6],[3,4,5,6]])

this works:

for i in range(m):
    R[i][S[i]]=C[i]

this does not work:

R[S]=C

You might get more responses, if you added a sample input and output to your description of the problem. — Mr. T, Jan 14 '18 at 20:15
Not sure but does `idx = np.ogrid[tuple(map(slice, R.shape))]` `idx[-1] = S` `R[idx] = C` do what you want? — Paul Panzer, Jan 14 '18 at 20:23
Please read [this section](https://stackoverflow.com/help/mcve). — Cleb, Jan 14 '18 at 20:25
A recent example using the index array from `argsort` on the last dimension of a 3d array: https://stackoverflow.com/questions/48072007/how-to-use-numpy-argsort-as-indices-in-more-than-2-dimensions — hpaulj, Jan 14 '18 at 21:20
@PaulPanzer, your code seems indeed to do what I want, but I do not understand how it works? Would you mind elaborating a litte bit? Thanks — Tashi Walde, Jan 14 '18 at 21:53
TashiWalde have a look at the link @hpaulj provided. My answer there tries to explain essentially the same thing. — Paul Panzer, Jan 14 '18 at 22:06
@PaulPanzer, thank you very much. If you post your remark as an answer I can accept it. — Tashi Walde, Jan 15 '18 at 11:43

score 0 · Accepted Answer · answered Jan 15 '18 at 18:08

Your 2D example can be done as follows:

R[np.arange(m)[:, None], S] = C
# array([[ 0.,  0.,  1.,  0.,  2.,  0.,  3.],
#        [ 0.,  0.,  0.,  4.,  5.,  6.,  7.]])

The 3D case would be similar:

i, j, k = R.shape
i, j, k = np.ogrid[i, j, k]
R[i, j, S] = C

In ND one could write:

idx = np.ogrid[tuple(map(slice, R.shape))]
idx[-1] = S
R[idx] = C

python: broadcast sliced np.array assignment to any number of dimensions

1 Answers1