CODEWITHSUNDEEP
numpy
user48956
user48956

Reputation: 15788

How to set cells of matrix from matrix of columns indexes

I'd like to build a kernel from a list of positions and list of kernel centers. The kernel should be an indicator of the TWO closest centers to each position.

> x = np.array([0.1, .49, 1.9, ]).reshape((3,1))  # Positions
> c = np.array([-2., 0.1,  0.2,  0.4,  0.5,  2.]) # centers
print x
print c

[[ 0.1 ]
 [ 0.49]
 [ 1.9 ]]
[-2.   0.1  0.2  0.4  0.5  2. ]

What I'd like to get out is: 

array([[ 0, 1, 1, 0, 0, 0],  # Index 1,2 closest to 0.1
       [ 0, 0, 0, 1, 1, 0],  # Index 3,4 closest to 0.49
       [ 0, 0, 0, 0, 1, 1]])  # Index 4,5 closest to  1.9

I can get:

> dist = np.abs(x-c)
array([[ 2.1 ,  0.  ,  0.1 ,  0.3 ,  0.4 ,  1.9 ],
       [ 2.49,  0.39,  0.29,  0.09,  0.01,  1.51],
       [ 3.9 ,  1.8 ,  1.7 ,  1.5 ,  1.4 ,  0.1 ]])

and:

> np.argsort(dist, axis=1)[:,:2]
array([[1, 2],
       [4, 3],
       [5, 4]])

Here I have a matrix of column indexes, but I but can't see how to use them to set values of those columns in another matrix (using efficient numpy operations).

idx = np.argsort(dist, axis=1)[:,:2]
z = np.zeros(dist.shape)
z[idx]=1 # NOPE
z[idx,:]=1 # NOPE
z[:,idx]=1 # NOPE

Upvotes: 0

Views: 36

Answers (1)

Divakar
Divakar

Reputation: 221574

One way would be to initialize zeros array and then index with advanced-indexing -

out = np.zeros(dist.shape,dtype=int)
out[np.arange(idx.shape[0])[:,None],idx] = 1

Alternatively, we could play around with dimensions extension to use broadcasting and come up with a one-liner -

out = (idx[...,None] == np.arange(dist.shape[1])).any(1).astype(int)

For performance, I would suggest using np.argpartition to get those indices - 

idx = np.argpartition(dist, 2, axis=1)[:,:2]

Upvotes: 1

Related Questions