Problem Statement in English

Given a 2D matrix grid, sort each upper triangular diagonal in ascending order, and each lower triangular diagonal in descending order.

Approach

Since they’ve asked us to deal with the upper and lower matrix triangles separately, we need to carefully control the iteration bounds.

Solution in Python


class Solution:
    def sortMatrix(self, grid: List[List[int]]) -> List[List[int]]:
        rows = len(grid)
        cols = len(grid[0])

        # upper triangle
        for c in range(1, cols):
            i = 0
            j = c
            maxheap = []

            while i < rows and j < cols:
                maxheap.append(grid[i][j])
                i += 1
                j += 1
            
            heapify(maxheap)
            i = 0
            j = c
            
            while maxheap:
                v = heappop(maxheap)
                grid[i][j] = v
                i += 1
                j += 1
        
        # lower triangle
        for r in range(rows):
            i = r
            j = 0
            maxheap = []

            while i < rows:
                maxheap.append(-grid[i][j])
                i += 1
                j += 1
            
            heapify(maxheap)
            i = r
            j = 0
            
            while maxheap:
                v = heappop(maxheap)
                grid[i][j] = -v
                i += 1
                j += 1

        return grid

Complexity

  • Time: $O(n^2log(n))$
    Since we go over each diagonal, and for each diagonal we sort the elements.

  • Space: $O(n)$
    For the max heap used to store the diagonal elements.

And we are done.