What is the Easiest way to extract subset of a 2D matrix in python?

Question

mat = [[0, 1, 2, 3, 4, 5],
         [6, 7, 8, 9, 10, 11],
         [12, 13, 14, 15, 16, 17],
         [18, 19, 20, 21, 22, 23],
         [24, 25, 26, 27, 28, 29],
         [30, 31, 32, 33, 34, 35]]

Lets say I want to extract upper left 2x2 matrix

[[0, 1,],
[6, 7, ]]

doing mat2=mat[:2][:2] doesnt work. It extracts the rows correctly but not columns.Seems like I need to loop throughto get the columns.

Additionally I need to do a deepcopy to mat2 suchthat modifying mat2 dont change mat.

Answer 1

This is because [:2] returns a list containing the first 2 elements of your matrix.

For example :-

arr = [[1, 2], [1, 3]]
print(arr[:2]) # will print the first 2 elements of the array, that is [1, 2] and [1, 3], packed into a list. So, Output : [[1, 2], [1, 3]].

In the same way,

mat = [[0, 1, 2, 3, 4, 5],
       [6, 7, 8, 9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23],
       [24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35]]

mat2 = mat[:2] # => [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]

# Now, if you again try to get the first 2 elements from mat2 you will get the first 2 elements of mat2, not the first 2 elements of the lists inside mat2.
mat3 = mat2[:2] # => [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]

That is where you went wrong, but this concept is quite counter-intuitive, so no worries.

---

So the solution would be to get the first 2 elements from matrix mat and then loop over its elements and then get the first 2 elements from them.

Therefore, this should work for you:

list(x[:2] for x in mat[:2])

Or, as @warped pointed, if you can use numpy, you can do the following:

import numpy as np

mat = np.array(mat)
mat[:2, :2]