Split a matrix in N chunks-CodePudding

Given a matrix, I want to split it in equally smaller matrices of m x n size. If the matrix is not divisible by the given size, we just put the remainder into a different matrix.

For example, given the matrix below and m=2 and n=2:

[[1, 0, 1],
 [0, 0, 0],
 [0, 1, 1]]

Result:

[[1, 0],
 [0, 0]],

[[1],
 [0]],

[[0, 1]],

[[1]],

I was using np.reshape but it fails to split when the numbers don't match, as in the example above.

    matrix_size = matrix.shape[0] * matrix.shape[1]
    n_matrix = math.ceil(matrix_size / (m * n))
    matrix.reshape(n_matrix, m, n)

CodePudding user response：

One way you could do this is using multiple calls to numpy.array_split

import numpy as np

matrix = [
    [1, 0, 1],
    [0, 0, 0],
    [0, 1, 1],
]
sub_matrices = np.array_split(matrix, 2, axis=0)
sub_matrices = [m for sub_matrix in sub_matrices for m in np.array_split(sub_matrix, 2, axis=1)]

Where the first call to array_split splits it vertically, and the second call splits it horizontally.