Runtime: 68 ms (69.15%), Memory: 17.1 MB (22.81%) - Gitcode

Author: vinija

1329-sort-the-matrix-diagonally/1329. Sort the Matrix Diagonally.py

@@ -0,0 +1,19 @@
+from collections import defaultdict
+import heapq
+from typing import List
+
+class Solution:
+    def diagonalSort(self, mat: List[List[int]]) -> List[List[int]]:
+        diagonals = defaultdict(list)
+        
+        # Collect all elements of each diagonal
+        for i in range(len(mat)):
+            for j in range(len(mat[0])):
+                heapq.heappush(diagonals[i - j], mat[i][j])
+        
+        # Put the sorted elements back into the matrix
+        for i in range(len(mat)):
+            for j in range(len(mat[0])):
+                mat[i][j] = heapq.heappop(diagonals[i - j])
+        
+        return mat