Day 23 - Sort the Matrix Diagonally

The Problem

A matrix diagonal is a diagonal line of cells starting from some cell in either the topmost row or leftmost column and going in the bottom-right direction until reaching the matrix's end. For example, the matrix diagonal starting from mat[2][0], where mat is a 6 x 3 matrix, includes cells mat[2][0], mat[3][1], and mat[4][2].

Given an m x n matrix mat of integers, sort each matrix diagonal in ascending order and return the resulting matrix.

Example 1:

Input: mat = [[3,3,1,1],[2,2,1,2],[1,1,1,2]]
Output: [[1,1,1,1],[1,2,2,2],[1,2,3,3]]

Constraints:

• m == mat.length
• n == mat[i].length
• 1 <= m, n <= 100
• 1 <= mat[i][j] <= 100

Tests

import pytest
from .Day23_SortMatrixDiagonally import Solution

s = Solution()


@pytest.mark.parametrize(
    "mat,expected",
    [
        (
            [[3, 3, 1, 1], [2, 2, 1, 2], [1, 1, 1, 2]],
            [[1, 1, 1, 1], [1, 2, 2, 2], [1, 2, 3, 3]],
        ),
    ],
)
def test_diagonal_sort(mat, expected):
    assert s.diagonalSort(mat) == expected

Solution

from typing import List


class Solution:
    def sort(self, mat, x, y):
        current_diag = []
        while x < len(mat) and y < len(mat[0]):
            current_diag.append(mat[x][y])
            x += 1
            y += 1
        current_diag.sort()
        while x > 0 and y > 0:
            x -= 1
            y -= 1
            mat[x][y] = current_diag.pop()

    def diagonalSort(self, mat: List[List[int]]) -> List[List[int]]:
        # For each diagonal in the matrix starting at the bottom left
        for i in range(len(mat)):
            self.sort(mat, i, 0)
        # For each diagonal in the matrix starting at the top right
        for i in range(len(mat[0])):
            self.sort(mat, 0, i)

        return mat

Analysis

This solution was accepted but the analysis view on leetcode appears to be broken currently.

Commentary

The simplest way I could imagine to solve this was to add each diagonal to a list, sort that list and then put the sorted list back in place of the diagonal.