aboutsummaryrefslogtreecommitdiff
path: root/sort_matrix_by_diagonals.py
blob: 1aee13e6064b5ed5b89b95c274a0ab909ebed806 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
# 3446. Sort Matrix by Diagonals
import heapq

"""
you are given an 'n x n' square matrix of integers 'grid'. return the matrix
such that the diagonals in the bottom left triangle including the middle are
sorted in non-increasing order. the diagonals in the top-right triangle are
sorted in non-decreasing order.
"""


class Solution(object):
    def sortMatrix(self, grid):
        """
        :type grid: List[List[int]]
        :rtype: List[List[int]]
        """
        n, m = len(grid), len(grid[0])
        diags = {}
        for i in range(n):
            for j in range(m):
                key = i - j
                if key not in diags:
                    diags[key] = []
                if key < 0:
                    heapq.heappush(diags[key], grid[i][j])
                else:
                    heapq.heappush(diags[key], -grid[i][j])
        for i in range(n):
            for j in range(m):
                key = i - j
                if key < 0:
                    grid[i][j] = heapq.heappop(diags[key])
                else:
                    grid[i][j] = -heapq.heappop(diags[key])
        return grid


if __name__ == "__main__":
    obj = Solution()
    print(obj.sortMatrix(grid=[[1, 7, 3], [9, 8, 2], [4, 5, 6]]))
    print(obj.sortMatrix(grid=[[0, 1], [1, 2]]))
    print(obj.sortMatrix(grid=[[1]]))