Search code examples
pythonnumpynumpy-ndarrayarray-broadcastingnumpy-slicing

Slice a submatrix with center element indices


Given a matrix A, and a list of row indices, and a list of column indices, how to efficiently extract the squared submatrices with size k centered by the row and column indices?

For example:

A = array([[12,  6, 14,  8,  4,  1],
       [18, 13,  8, 10,  9, 19],
       [ 8, 15,  6,  5,  6, 18],
       [ 3,  0,  2, 14, 13, 12],
       [ 4,  4,  5, 19,  0, 14],
       [16,  8,  7,  7, 11,  0],
       [ 3, 11,  2, 19, 11,  5],
       [ 4,  2,  1,  9, 12, 12]])
r = np.array([2, 5])
c = np.array([3, 2])
k = 3

The output should be A[1:4, 2:5] and A[4:7, 1:4]. So basically, the outputs are squared submatrices in size kxk and centered on the [r,c] elements (A[2,3] and A[5,2] in this case)

How to do this efficiently and elegantly? Thanks


Solution

  • For the case when the submatrices be of the same shape, we can get sliding windows and then index those with the start indices along the rows and cols for our desired output. To get those windows, we can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows. More info on use of as_strided based view_as_windows -

    from skimage.util.shape import view_as_windows
    
    # Get all sliding windows
    w = view_as_windows(A,(k,k))
    
    # Select relevant ones for final o/p
    out = w[r-k//2,c-k//2]