Python | Get Kth Column of Matrix

Sometimes, while working with Python Matrix, one can have a problem in which one needs to find the Kth column of Matrix. This is a very popular problem in Machine Learning Domain and having solution to this is useful. Let’s discuss certain ways in which this problem can be solved. 

Method #1 : Using list comprehension This problem can be solved using list comprehension in which we can iterate through all the rows and selectively gather all the elements occurring at Kth index. 

Python3

# Python3 code to demonstrate working of
# Get Kth Column of Matrix
# using list comprehension

# initialize list
test_list = [[4, 5, 6], [8, 1, 10], [7, 12, 5]]

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

# initialize K
K = 2

# Get Kth Column of Matrix
# using list comprehension
res = [sub[K] for sub in test_list]

# printing result
print("The Kth column of matrix is : " + str(res))

Output : 

The original list is : [[4, 5, 6], [8, 1, 10], [7, 12, 5]]
The Kth column of matrix is : [6, 10, 5]

Time complexity: O(n), where n is the number of elements in the matrix.
Auxiliary space: O(1), as only constant space is used to store the input matrix and the integer K.

Method #2 : Using zip() This task can also be performed using zip(). This does the similar task of gathering elements like is done by above list comprehension and offers compact but slower execution. Works with Python2 only. 

Python

# Python code to demonstrate working of
# Get Kth Column of Matrix
# using zip()

# initialize list
test_list = [[4, 5, 6], [8, 1, 10], [7, 12, 5]]

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

# initialize K
K = 2

# Get Kth Column of Matrix
# using zip()
res = list(zip(*test_list)[K])

# printing result
print("The Kth column of matrix is : " + str(res))

Output : 

The original list is : [[4, 5, 6], [8, 1, 10], [7, 12, 5]]
The Kth column of matrix is : [6, 10, 5]

Time Complexity: O(n), where n is the length of the list test_list 
Auxiliary Space: O(n) additional space of size n is created where n is the number of elements in the res list 

Method #3: Using numpy

Note: First Install numpy module using command “pip install numpy”

This problem can also be solved using numpy library which provides a very efficient and easy way to get the Kth column of matrix.

Python3

import numpy as np
# initialize list
test_list = [[4, 5, 6], [8, 1, 10], [7, 12, 5]]
# printing original list
print("The original list is : " + str(test_list))
#initialize K
K = 2
# Get Kth Column of Matrix
# using numpy
res = np.array(test_list)[:,K]
# printing result
print("The Kth column of matrix is : " + str(res))
#This code is contributed by Edula Vinay Kumar Reddy

Output:

The original list is : [[4, 5, 6], [8, 1, 10], [7, 12, 5]]
The Kth column of matrix is : [ 6 10  5]

Time Complexity: O(n*n), where n is the length of the list test_list 
Auxiliary Space: O(n) where n is the length of the list test_list 

Method 4: Using a for loop.

This code uses a for loop to iterate over the rows of the matrix and extract the Kth element from each row using indexing. The extracted elements are then appended to a new list called res. The final output is the list res containing the Kth column of the matrix.

Python3

test_list = [[4, 5, 6], [8, 1, 10], [7, 12, 5]]

K = 2

res = []

for i in range(len(test_list)):
    res.append(test_list[i][K])

print("The Kth column of matrix is : " + str(res))

Output

The Kth column of matrix is : [6, 10, 5]

Time complexity: O(n), where n is the number of rows in the matrix.
Auxiliary space:  O(n), because the list “res” will have at most n elements.

Method #6: Using the built-in function map()

Use the map() function to apply a lambda function that extracts the Kth element from each sublist in the matrix. The resulting map object is then converted to a list to obtain the final result.

Python3

# initialize list
test_list = [[4, 5, 6], [8, 1, 10], [7, 12, 5]]

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

# initialize K
K = 2

# Get Kth Column of Matrix using map()
res = list(map(lambda x: x[K], test_list))

# printing result
print("The Kth column of matrix is : " + str(res))

Output

The original list is : [[4, 5, 6], [8, 1, 10], [7, 12, 5]]
The Kth column of matrix is : [6, 10, 5]

The time complexity of this code is O(n), where n is the total number of elements in the input matrix.
The space complexity of this code is O(n), where n is the total number of elements in the input matrix. 

Method 6: Using list slicing

Step-by-step approach:

1. Initialize the list to store the kth column of the matrix
2. Loop through each row in the matrix
3. Append the element at index k of the row to the list initialized in step 1
4. Return the list as the kth column of the matrix

Below is the implementation of the above approach:

Python3

# Python code to demonstrate working of
# Get Kth Column of Matrix
# using list slicing

# initialize list
test_list = [[4, 5, 6], [8, 1, 10], [7, 12, 5]]

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

# initialize K
K = 2

# Get Kth Column of Matrix
# using list slicing
res = [row[K] for row in test_list]

# printing result
print("The Kth column of matrix is : " + str(res))

Output

The original list is : [[4, 5, 6], [8, 1, 10], [7, 12, 5]]
The Kth column of matrix is : [6, 10, 5]

Time complexity: O(n), where n is the number of elements in the matrix.
Auxiliary space: O(n), where n is the number of elements in the matrix.