David Chen
David Chen

Reputation: 1835

delete all columns of a dimension except for a specific column

I want to make a function which takes a n-dimensional array, the dimension and the column index, and it will return the (n-1)-dimensional array after removing all the other columns of that specific dimension.

Here is the code I am using now

a = np.arange(6).reshape((2, 3))  # the n-dimensional array
axisApplied = 1
colToKeep = 0
colsToDelete = np.delete(np.arange(a.shape[axisApplied]), colToKeep)
a = np.squeeze(np.delete(a, colsToDelete, axisApplied), axis=axisApplied)
print(a)
# [0, 3]

Note that I have to manually calculate the n-1 indices (the complement of the specific column index) to use np.delete(), and I am wondering whether there is a more convenient way to achieve my goal, e.g. specify which column to keep directly.

Thank you for reading and I am welcome to any suggestions.

Upvotes: 0

Views: 342

Answers (1)

hpaulj
hpaulj

Reputation: 231475

In [1]: arr = np.arange(6).reshape(2,3)
In [2]: arr
Out[2]: 
array([[0, 1, 2],
       [3, 4, 5]])

Simple indexing:

In [3]: arr[:,0]
Out[3]: array([0, 3])

Or if you need to used the general axis parameter, try take:

In [4]: np.take(arr,0,axis=1)
Out[4]: array([0, 3])

Picking one element, or a list of elements, along an axis is a lot easier than deleting some. Look at the code for np.delete.

Upvotes: 2

Related Questions