Python | Get Kth Column of Matrix

• Difficulty Level : Easy
• Last Updated : 30 Dec, 2020

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 code to demonstrate working of# Get Kth Column of Matrix# using list comprehension  # initialize listtest_list = [[4, 5, 6], [8, 1, 10], [7, 12, 5]]  # printing original listprint("The original list is : " + str(test_list))  # initialize KK = 2  # Get Kth Column of Matrix# using list comprehensionres = [sub[K] for sub in test_list]  # printing resultprint("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]

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 code to demonstrate working of# Get Kth Column of Matrix# using zip()  # initialize listtest_list = [[4, 5, 6], [8, 1, 10], [7, 12, 5]]  # printing original listprint("The original list is : " + str(test_list))  # initialize KK = 2  # Get Kth Column of Matrix# using zip()res = list(zip(*test_list)[K])  # printing resultprint("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]

My Personal Notes arrow_drop_up