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.

filter_none

edit
close

play_arrow

link
brightness_4
code

# 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))

chevron_right


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.

filter_none

edit
close

play_arrow

link
brightness_4
code

# 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))

chevron_right


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

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.




Article Tags :

Be the First to upvote.


Please write to us at contribute@geeksforgeeks.org to report any issue with the above content.