CODEWITHSUNDEEP
pythonlistmatrixslice
clouvis
clouvis

Reputation: 575

Python: slice a matrix horizontally

Let's consider those two variables :

matrix = [[1, 2, 3, 4],
          [5, 6, 7, 8],
          [9, 10, 11, 12]]

list_of_lengths = [2, 1, 1]

I am looking for a method to get this result :

result = [[[1, 2],
           [5, 6],
           [9, 10]], [[3],
                      [7],
                      [11]], [[4],
                              [8],
                              [12]]]

Indeed, I want three matrix as a result. The j-length of each one is determined by the variable list_of_lentghts.

I have already written a solution that works, but I would like to know if there is a solution more simple. Here is my solution:

def cut_matrix(matrix, list_of_lengths):
    result = []
    for a_len in list_of_lengths:
        current_sub_matrix = []
        for a_line in matrix:
            current_new_line = []
            for i in range(0, a_len):
                current_new_line.append(a_line.pop(0))
            current_sub_matrix.append(current_new_line)
        result.append(current_sub_matrix)
    return result

Upvotes: 2

Views: 519

Answers (1)

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 476967

If you know the column offset i and the number of columns n, you can use slicing to obtain the columns:

[row[i:i+n] for row in matrix]

to get the slice. So you can simply use:

def cut_matrix(matrix, list_of_lengths):
    result = []
    col = 0
    for a_len in list_of_lengths:
        result.append([row[col:col+a_len] for row in matrix])
        col += a_len
    return result

It generates:

>>> cut_matrix(matrix,list_of_lengths)
[[[1, 2], [5, 6], [9, 10]], [[3], [7], [11]], [[4], [8], [12]]]

This will also work faster than using .pop(0) since popping from the front is done in O(n) (so popping all elements requires O(n2) whereas slicing all elements is done in O(n)). Finally it leaves the original matrix intact which is thus more declarative.

Upvotes: 2

Related Questions