Question

I want to extract out some specified columns from the matrix. My matrix being

matrix=[[1,2,3,4,5],
        [6,7,8,9,10],
        [6,4,3,1,2],
        [2,3,4,5,6]]
expected result: [[3,4],[8,9],[3,1],[4,5]] for 2 columns
expected result in case of 3 column:[[1,2,3],[6,7,8],[6,4,3],[2,3,4]]

I am trying to use the method given below:

def func(matrix, p):
     return np.vstack([a[i] for a in matrix])

The above method just returns to be a single column, however I want to write a method which takes in multiple columns as input (e.g. the expected results for 2 & 3 columns), also my input number of columns vary each time. Please suggest a suitable method to extract these columns in python.

Answer 1

Using pure Python, you can extract columns 1, 5 and 7 by using the following nested list comprehension:

[[a[i] for i in (1, 5, 7)] for a in matrix]

Example usage for your input:

>>> [[a[i] for i in (0, 1, 2)] for a in matrix]
[[1, 2, 3], [6, 7, 8], [6, 4, 3], [2, 3, 4]]

In cases like this where you want to extract a real slice of columns, you can also use list slicing:

>>> [a[0:3] for a in matrix]
[[1, 2, 3], [6, 7, 8], [6, 4, 3], [2, 3, 4]]