aboutsummaryrefslogtreecommitdiff
path: root/sort_matrix_by_diagonals.py
diff options
context:
space:
mode:
Diffstat (limited to 'sort_matrix_by_diagonals.py')
-rw-r--r--sort_matrix_by_diagonals.py43
1 files changed, 43 insertions, 0 deletions
diff --git a/sort_matrix_by_diagonals.py b/sort_matrix_by_diagonals.py
new file mode 100644
index 0000000..1aee13e
--- /dev/null
+++ b/sort_matrix_by_diagonals.py
@@ -0,0 +1,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]]))