CODEWITHSUNDEEP
pythonmatplotlibscipypython-2.7
Jannat Arora
Jannat Arora

Reputation: 2989

Extract variable number of columns for a matrix

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.

Upvotes: 1

Views: 535

Answers (2)

Niklas B.
Niklas B.

Reputation: 95308

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]]

Upvotes: 3

mechanical_meat
mechanical_meat

Reputation: 169334

If you prefer to work within Numpy you could use slicing:

def col_slice(arr, col_start=None, total_cols=None):
    """ take a slice of columns from col_start for total_cols """
    return arr[:,col_start:col_start + total_cols]

Examples:

In [22]: col_slice(your_matrix, col_start=2, total_cols=2)
Out[22]:
array([[3, 4],
       [8, 9],
       [3, 1],
       [4, 5]])

In [23]: col_slice(your_matrix, col_start=0, total_cols=3)
Out[23]:
array([[1, 2, 3],
       [6, 7, 8],
       [6, 4, 3],
       [2, 3, 4]])

Upvotes: 1

Related Questions