Wychh
Wychh

Reputation: 713

Print multiple columns from a matrix

I have a list of column vectors and I want to print only those column vectors from a matrix. Note: the list can be of random length, and the indices can also be random.

For instance, the following does what I want:

import numpy as np

column_list = [2,3]
a = np.array([[1,2,6,1],[4,5,8,2],[8,3,5,3],[6,5,4,4],[5,2,8,8]])

new_matrix = []
for i in column_list:
    new_matrix.append(a[:,i])
new_matrix = np.array(new_matrix)
new_matrix = new_matrix.transpose()

print(new_matrix)

However, I was wondering if there is a shorter method?

Upvotes: 0

Views: 338

Answers (1)

user17242583
user17242583

Reputation:

Yes, there's a shorter way. You can pass a list (or numpy array) to an array's indexer. Therefore, you can pass column_list to the columns indexer of a:

>>> a[:, column_list]
array([[6, 1],
       [8, 2],
       [5, 3],
       [4, 4],
       [8, 8]])
       
# This is your new_matrix produced by your original code:
>>> new_matrix
array([[6, 1],
       [8, 2],
       [5, 3],
       [4, 4],
       [8, 8]])
       
>>> np.all(a[:, column_list] == new_matrix)
True

Upvotes: 1

Related Questions