get columns from numpy array in specific format given column indices

Question

I have a numpy array like this:

import numpy as np

# Create a NumPy array
array_data = np.array([[1, 2, 3],
                       [4, 5, 6],
                       [7, 8, 9]])

and yes I can get the 2nd column like this:

column_index = 1
selected_column = array_data[:,1]
print(selected_column)

giving:

[2 5 8]

I need it like this though:

[[2]
 [5]
 [8]]

which I can get using:

selected_column = array_data[:, column_index:column_index+1]
print(selected_column)

Now you wonder, why ask this question. Well I would like to actually provide several column indices (coming from a pandas data frame with provided column names) whilst still getting the array like the last output.

selected_column = array_data[:, column_index:column_index+1]

Asked By: cs0815

||

Source

Answer 1

Try the following:

selected_column_vector = array_data[:, [1]]

Which can be used to fetch several columns too:

column_index = 1
result = array_data[:, [column_index, column_index + 1]]

Answered By: Guru Stron

Answer 2

You could use reshape?

array_data[:, column_index].reshape(-1, 1)

Answered By: SomeDude

get columns from numpy array in specific format given column indices

Question:

Answers: