Assign values to a numpy array for each row with specified columns

Question

I have a matrix foo with n rows and m columns. Example:

>>> import numpy as np
>>> foo = np.arange(6).reshape(3, 2) # n=3 and m=2 in our example
>>> print(foo)
array([[0, 1],
       [2, 3],
       [4, 5]])

I have an array bar with n elements. Example:

>>> bar = np.array([9, 8, 7])

I have a list ind of length n that contains column indices. Example:

>>> ind = np.array([0, 0, 1], dtype='i')

I would like to use the column indices ind to assign the values of bar to the matrix foo. I would like to do this per row. Assume that the function that does this is called assign_function, my output would look as follows:

>>> assign_function(ind, bar, foo)
>>> print(foo)
array([[9, 1],
       [8, 3],
       [4, 7]])

Is there a pythonic way to do this?

To start with, the result you say to with to be getting is not corresponding to the variables defined, right? — user4396006
– user4396006, Commented Dec 19, 2017 at 16:31

cs95 · Accepted Answer · 2017-12-19 16:30:42Z

7

Since ind takes care of the first axis, you just need the indexer for the zeroth axis. You can do this pretty simply with np.arange:

foo[np.arange(len(foo)), ind] = bar
foo

array([[9, 1],
       [8, 3],
       [4, 7]])

answered Dec 19, 2017 at 16:30

cs95

406k106 gold badges744 silver badges797 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Divakar · Accepted Answer · 2017-12-19 16:31:52Z

3

Leveraging broadcasting alongwith masking -

foo[ind[:,None] == range(foo.shape[1])] = bar

Sample step-by-step run -

# Input array
In [118]: foo
Out[118]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

# Mask of places to be assigned
In [119]: ind[:,None] == range(foo.shape[1])
Out[119]: 
array([[ True, False],
       [ True, False],
       [False,  True]], dtype=bool)

# Assign values off bar
In [120]: foo[ind[:,None] == range(foo.shape[1])] = bar

# Verify
In [121]: foo
Out[121]: 
array([[9, 1],
       [8, 3],
       [4, 7]])

answered Dec 19, 2017 at 16:31

Divakar

222k19 gold badges273 silver badges374 bronze badges

Collectives™ on Stack Overflow

Assign values to a numpy array for each row with specified columns

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related