CODEWITHSUNDEEP
pythonnumpypandas
Dickster
Dickster

Reputation: 3009

numpy - update values using slicing given an array value

Assume I have the following array

import numpy as np
a = np.arange(0,36).reshape((6,6)).T

[[ 0  6 12 18 24 30]
 [ 1  7 13 19 25 31]
 [ 2  8 14 20 26 32]
 [ 3  9 15 21 27 33]
 [ 4 10 16 22 28 34]
 [ 5 11 17 23 29 35]]


for i in a[:,0]:
    a[i][i:] = 0

[[ 0  0  0  0  0  0]
 [ 1  0  0  0  0  0]
 [ 2  8  0  0  0  0]
 [ 3  9 15  0  0  0]
 [ 4 10 16 22  0  0]
 [ 5 11 17 23 29  0]]

I wish to know if I can update (wipe out) the values using the "first column" as an indicator of the start of the slice on axis =1 and do this without using a loop.

Please note the values in "first column" will not necessarily be in order as shown in the example so numpy.tril is not appropriate for me here. I do know that the value in "first column" will never be bigger than the size of axis=1.

Upvotes: 1

Views: 1073

Answers (1)

DSM
DSM

Reputation: 353589

How about something like this? Note that I've shuffled the first column.

>>> a = np.arange(0,36).reshape((6,6)).T; a[2,0] = 4; a[4,0] = 2;
>>> a
array([[ 0,  6, 12, 18, 24, 30],
       [ 1,  7, 13, 19, 25, 31],
       [ 4,  8, 14, 20, 26, 32],
       [ 3,  9, 15, 21, 27, 33],
       [ 2, 10, 16, 22, 28, 34],
       [ 5, 11, 17, 23, 29, 35]])
>>> a[np.arange(a.shape[1])[None,:] >= a[:,0,None]] = 0
>>> a
array([[ 0,  0,  0,  0,  0,  0],
       [ 1,  0,  0,  0,  0,  0],
       [ 4,  8, 14, 20,  0,  0],
       [ 3,  9, 15,  0,  0,  0],
       [ 2, 10,  0,  0,  0,  0],
       [ 5, 11, 17, 23, 29,  0]])

Upvotes: 1

Related Questions