Python - Custom Columns Matrix

Sometimes, while working with Python lists, we can have a problem in which we need to extract certain columns from Matrix and recreate it. This kind of problem can have applications in data domains as they use Matrix as a prominent input parameter. Let's discuss certain ways in which this task can be performed.

Input : test_list = [[5, 4, 3, 4], [7, 6, 3, 2], [8, 3, 9, 10]], col_list = [2] 
Output : [[3], [3], [9]] 

Input : test_list = [[5, 4], [6, 2], [8, 3]], col_list = [1] 
Output : [[4], [2], [3]]

Method #1: Using list comprehension This offers one of the ways to solve this problem. In this, we perform extraction of selective columns using nested list comprehension. 

# Python3 code to demonstrate working of 
# Custom Columns Matrix
# Using list comprehension

# initializing list
test_list = [[5, 4, 3, 4],
             [7, 6, 3, 2], 
             [8, 3, 9, 10]]

# printing original list 
print("The original list : " + str(test_list))

# initializing Columns list 
col_list = [1, 3]

# Custom Columns Matrix
# Using list comprehension
res = [[sub[idx] for idx in col_list] for sub in test_list]

# printing result 
print("Matrix after filtering : " + str(res))

The original list : [[5, 4, 3, 4], [7, 6, 3, 2], [8, 3, 9, 10]]
Matrix after filtering : [[4, 4], [6, 2], [3, 10]]

Time complexity: O(m*n), where m is the number of rows and n is the number of columns in the input matrix.
Auxiliary space: O(k), where k is the length of the column list. This is because we are creating a new matrix to store the filtered values, which will have k columns.

Method #2: Using itemgetter() + list comprehension The combination of above functions can be used to solve this problem. In this, we perform the task of getting indices using itemgetter(). 

# Python3 code to demonstrate working of 
# Custom Columns Matrix
# Using itemgetter() + list comprehension
from operator import itemgetter

# initializing list
test_list = [[5, 4, 3, 4],
             [7, 6, 3, 2], 
             [8, 3, 9, 10]]

# printing original list 
print("The original list : " + str(test_list))

# initializing Columns list 
col_list = [1, 3]

# Custom Columns Matrix
# Using itemgetter() + list comprehension
res = [list(itemgetter(*col_list)(ele)) for ele in test_list]

# printing result 
print("Matrix after filtering : " + str(res))

The original list : [[5, 4, 3, 4], [7, 6, 3, 2], [8, 3, 9, 10]]
Matrix after filtering : [[4, 4], [6, 2], [3, 10]]

Time complexity: O(nm), where n is the number of rows in the input list and m is the number of columns. This is because the program loops through each row of the input list and then selects the specified columns using the itemgetter() function, which takes constant time.
Auxiliary space: O(nm), where n is the number of rows and m is the number of columns, since the program creates a new matrix of the same size as the input matrix to store the selected columns.

Method #3: Using a nested for loop to iterate over the elements and columns of the original list

Initializes an empty result list, then iterates over each element of the original list. For each element, a new empty row list is created. The code then iterates over the columns specified in col_list and appends the corresponding values to the row list. Once all columns have been iterated over, the row list is appended to the result list. The resulting matrix is then printed.

# Python3 code to demonstrate working of 
# Custom Columns Matrix
# Using nested for loop

# initializing list
test_list = [[5, 4, 3, 4],
             [7, 6, 3, 2], 
             [8, 3, 9, 10]]

# printing original list 
print("The original list : " + str(test_list))

# initializing Columns list 
col_list = [1, 3]

# Custom Columns Matrix
# Using nested for loop
res = []
for ele in test_list:
    row = []
    for i in col_list:
        row.append(ele[i])
    res.append(row)

# printing result 
print("Matrix after filtering : " + str(res))

The original list : [[5, 4, 3, 4], [7, 6, 3, 2], [8, 3, 9, 10]]
Matrix after filtering : [[4, 4], [6, 2], [3, 10]]

Time complexity: O(n * m), where n is the number of rows and m is the number of columns in the original list. 
Auxiliary space: O(n * m), as the code creates a new list res to store the filtered matrix. 

Method 5: Using numpy indexing. 

import numpy as np

# initializing list
test_list = [[5, 4, 3, 4],
             [7, 6, 3, 2], 
             [8, 3, 9, 10]]

# initializing Columns list 
col_list = [1, 3]

# Custom Columns Matrix using numpy indexing
arr = np.array(test_list)
res = arr[:, col_list]

# printing original list 
print("The original list : " + str(test_list))

# printing result 
print("Matrix after filtering : " + str(res))

Output:

The original list : [[5, 4, 3, 4], [7, 6, 3, 2], [8, 3, 9, 10]]
Matrix after filtering : [[ 4  4]
[ 6  2]
[ 3 10]]

Time complexity: O(n), where n is the number of elements in the input list. 
Auxiliary space: O(n), where n is the number of elements in the input list, as we are creating a new numpy array to store the filtered columns.