CODEWITHSUNDEEP
pythonnumpy
user1610952
user1610952

Reputation: 1289

Delete one element from each row of a NumPy array

import numpy as np
a=np.array([[1,2,3], [4,5,6], [7,8,9]])

k = [0, 1, 2]

print np.delete(a, k, 1)

This returns

[]

But, the result I really want is

[[2,3],
 [4,6],
 [7,8]]

I want to delete the first element (indexed as 0) from a[0], the second (indexed as 1) from a[1], and the third (indexed as 2) from a[2].

Any thoughts?

Upvotes: 4

Views: 2202

Answers (1)

Divakar
Divakar

Reputation: 221714

Here's an approach using boolean indexing -

m,n = a.shape
out = a[np.arange(n) != np.array(k)[:,None]].reshape(m,-1)

If you would like to persist with np.delete, you could calculate the linear indices and then delete those after flattening the input array, like so -

m,n = a.shape
del_idx = np.arange(n)*m + k
out = np.delete(a.ravel(),del_idx,axis=0).reshape(m,-1)

Sample run -

In [94]: a
Out[94]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [95]: k = [0, 2, 1]

In [96]: m,n = a.shape

In [97]: a[np.arange(n) != np.array(k)[:,None]].reshape(m,-1)
Out[97]: 
array([[2, 3],
       [4, 5],
       [7, 9]])

In [98]: del_idx = np.arange(n)*m + k

In [99]: np.delete(a.ravel(),del_idx,axis=0).reshape(m,-1)
Out[99]: 
array([[2, 3],
       [4, 5],
       [7, 9]])

Upvotes: 4

Related Questions