Wendy
Wendy

Reputation: 871

Modify multiple columns in an array numpy

I have a numpy array (nxn matrix), and I would like to modify only the columns which sum is 0. And I would like to assign the same value to all of these columns. To do that, I have first taken the index of the columns that sum to 0:

sum_lines = np.sum(mat_trans, axis = 0)
indices = np.where(sum_lines == 0)[0]

then I did a loop on those indices:

for i in indices:
    mat_trans[:, i] = rank_vect

so that each of these columns now has the value of the rank_vect column vector.

I was wondering if there was a way to do this without loop, something that would look like:

mat_trans[:, (np.where(sum_lines == 0)[0]))] = rank_vect

Thanks!

Upvotes: 0

Views: 1805

Answers (2)

hpaulj
hpaulj

Reputation: 231375

In [114]: arr = np.array([[0,1,2,3],[1,0,2,-3],[-1,2,0,0]])

In [115]: sumlines = np.sum(arr, axis=0)
In [116]: sumlines
Out[116]: array([0, 3, 4, 0])

In [117]: idx = np.where(sumlines==0)[0]
In [118]: idx
Out[118]: array([0, 3])

So the columns that we want to modify are:

In [119]: arr[:,idx]
Out[119]: 
array([[ 0,  3],
       [ 1, -3],
       [-1,  0]])

In [120]: rv = np.array([10,11,12])

If rv is 1d, we get a shape error:

In [121]: arr[:,idx] = rv
ValueError: shape mismatch: value array of shape (3,) could not be broadcast to indexing result of shape (2,3)

But if it is a column vector (shape (3,1)) it can be broadcast to the (3,2) target:

In [122]: arr[:,idx] = rv[:,None]

In [123]: arr
Out[123]: 
array([[10,  1,  2, 10],
       [11,  0,  2, 11],
       [12,  2,  0, 12]])

Upvotes: 2

kevinkayaks
kevinkayaks

Reputation: 2726

This should do the trick

mat_trans[:,indices] = np.stack((rank_vect,)*indices.size,-1)

Please test and let me know if it does what you want. It just stacks the rank_vect repeatedly to match the shape of the LHS on the RHS.

I believe this is equivalent to

for i in indices:
    mat_trans[:, i] = rank_vec

I'd be interested to know the speed difference

Upvotes: 0

Related Questions