Sort matrix diagonally using heaps. (keon#665)

README.md

@@ -244,6 +244,7 @@ If you want to uninstall algorithms, it is as simple as:
 	- [crout_matrix_decomposition](algorithms/matrix/crout_matrix_decomposition.py)
 	- [cholesky_matrix_decomposition](algorithms/matrix/cholesky_matrix_decomposition.py)
     - [sum_sub_squares](algorithms/matrix/sum_sub_squares.py)
+    - [sort_matrix_diagonally](algorithms/matrix/sort_matrix_diagonally.py)
 - [queues](algorithms/queues)
     - [max_sliding_window](algorithms/queues/max_sliding_window.py)
     - [moving_average](algorithms/queues/moving_average.py)

algorithms/matrix/sort_matrix_diagonally.py

@@ -0,0 +1,77 @@
+"""
+Given a m * n matrix mat of integers,
+sort it diagonally in ascending order
+from the top-left to the bottom-right
+then return the sorted array.
+
+mat = [
+    [3,3,1,1],
+    [2,2,1,2],
+    [1,1,1,2]
+]
+
+Should return:
+[
+    [1,1,1,1],
+    [1,2,2,2],
+    [1,2,3,3]
+]
+"""
+
+from heapq import heappush, heappop
+from typing import List
+
+
+def sort_diagonally(mat: List[List[int]]) -> List[List[int]]:
+    # If the input is a vector, return the vector
+    if len(mat) == 1 or len(mat[0]) == 1:
+        return mat
+
+    # Rows + columns - 1
+    # The -1 helps you to not repeat a column
+    for i in range(len(mat)+len(mat[0])-1):
+        # Process the rows
+        if i+1 < len(mat):
+            # Initialize heap, set row and column
+            h = []
+            row = len(mat)-(i+1)
+            col = 0
+
+            # Traverse diagonally, and add the values to the heap
+            while row < len(mat):
+                heappush(h, (mat[row][col]))
+                row += 1
+                col += 1
+
+            # Sort the diagonal
+            row = len(mat)-(i+1)
+            col = 0
+            while h:
+                ele = heappop(h)
+                mat[row][col] = ele
+                row += 1
+                col += 1
+        else:
+            # Process the columns
+            # Initialize heap, row and column
+            h = []
+            row = 0
+            col = i - (len(mat)-1)
+
+            # Traverse Diagonally
+            while col < len(mat[0]) and row < len(mat):
+                heappush(h, (mat[row][col]))
+                row += 1
+                col += 1
+
+            # Sort the diagonal
+            row = 0
+            col = i - (len(mat)-1)
+            while h:
+                ele = heappop(h)
+                mat[row][col] = ele
+                row += 1
+                col += 1
+
+    # Return the updated matrix
+    return mat

tests/test_matrix.py

@@ -10,32 +10,33 @@
     sparse_dot_vector,
     spiral_traversal,
     sudoku_validator,
-    sum_sub_squares
+    sum_sub_squares,
+    sort_matrix_diagonally
 )
 import unittest
 
 
 class TestBombEnemy(unittest.TestCase):
     def test_3x4(self):
         grid1 = [
-                 ["0","E","0","0"],
-                 ["E","0","W","E"],
-                 ["0","E","0","0"]
-                ]
+            ["0", "E", "0", "0"],
+            ["E", "0", "W", "E"],
+            ["0", "E", "0", "0"]
+        ]
         self.assertEqual(3, bomb_enemy.max_killed_enemies(grid1))
 
         grid1 = [
-                 ["0", "E", "0", "E"],
-                 ["E", "E", "E", "0"],
-                 ["E", "0", "W", "E"],
-                 ["0", "E", "0", "0"]
-                ]
+            ["0", "E", "0", "E"],
+            ["E", "E", "E", "0"],
+            ["E", "0", "W", "E"],
+            ["0", "E", "0", "0"]
+        ]
         grid2 = [
-                 ["0", "0", "0", "E"],
-                 ["E", "0", "0", "0"],
-                 ["E", "0", "W", "E"],
-                 ["0", "E", "0", "0"]
-                ]
+            ["0", "0", "0", "E"],
+            ["E", "0", "0", "0"],
+            ["E", "0", "W", "E"],
+            ["0", "E", "0", "0"]
+        ]
         self.assertEqual(5, bomb_enemy.max_killed_enemies(grid1))
         self.assertEqual(3, bomb_enemy.max_killed_enemies(grid2))
 
@@ -112,7 +113,7 @@ def test_cholesky_matrix_decomposition(self):
                          cholesky_matrix_decomposition.cholesky_decomposition(
                              [[4, 12, -16], [12, 37, -43], [-16, -43, 98]]))
 
-        self.assertEqual( None,
+        self.assertEqual(None,
                          cholesky_matrix_decomposition.cholesky_decomposition(
                              [[4, 12, -8], [12, 4, -43], [-16, -1, 32]]))
 
@@ -123,24 +124,28 @@ def test_cholesky_matrix_decomposition(self):
         # example taken from https://ece.uwaterloo.ca/~dwharder/NumericalAnalysis/04LinearAlgebra/cholesky/
         self.assertEqual([[2.23606797749979, 0.0, 0.0, 0.0],
                           [0.5366563145999494, 2.389979079406345, 0.0, 0.0],
-                          [0.13416407864998736, -0.19749126846635062, 2.818332343581848, 0.0],
+                          [0.13416407864998736, -0.19749126846635062,
+                              2.818332343581848, 0.0],
                           [-0.2683281572999747, 0.43682390737048743, 0.64657701271919, 3.052723872310221]],
                          cholesky_matrix_decomposition.cholesky_decomposition(
                              [[5, 1.2, 0.3, -0.6], [1.2, 6, -0.4, 0.9],
                               [0.3, -0.4, 8, 1.7], [-0.6, 0.9, 1.7, 10]]))
 
+
 class TestInversion(unittest.TestCase):
     """[summary]
     Test for the file matrix_inversion.py
 
     Arguments:
         unittest {[type]} -- [description]
     """
+
     def test_inversion(self):
         from fractions import Fraction
 
         m1 = [[1, 1], [1, 2]]
-        self.assertEqual(matrix_inversion.invert_matrix(m1), [[2, -1], [-1, 1]])
+        self.assertEqual(matrix_inversion.invert_matrix(m1),
+                         [[2, -1], [-1, 1]])
 
         m2 = [[1, 2], [3, 4, 5]]
         self.assertEqual(matrix_inversion.invert_matrix(m2), [[-1]])
@@ -151,15 +156,17 @@ def test_inversion(self):
         m4 = [[1]]
         self.assertEqual(matrix_inversion.invert_matrix(m4), [[-3]])
 
-        m5 = [[1, 2, 3]   , [4, 5, 6], [7, 8, 9]]
+        m5 = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
         self.assertEqual(matrix_inversion.invert_matrix(m5), [[-4]])
 
         m6 = [[3, 5, 1], [2, 5, 0], [1, 9, 8]]
         self.assertEqual(matrix_inversion.invert_matrix(m6), [[Fraction(40, 53), Fraction(-31, 53), Fraction(-5, 53)],
-                                                              [Fraction(-16, 53), Fraction(23, 53), Fraction(2, 53)],
+                                                              [Fraction(-16, 53), Fraction(
+                                                                  23, 53), Fraction(2, 53)],
                                                               [Fraction(13, 53), Fraction(-22, 53), Fraction(5, 53)]])
 
 
+
 class TestMatrixExponentiation(unittest.TestCase):
     """[summary]
     Test for the file matrix_exponentiation.py
@@ -329,21 +336,39 @@ def test_sudoku_validator(self):
                     [3, 0, 0, 4, 8, 1, 1, 7, 9]
                 ]))
 
+
 class TestSumSubSquares(unittest.TestCase):
     """[summary]
     Test for the file sum_sub_squares.py
 
     Arguments:
         unittest {[type]} -- [description]
     """
+
     def test_sum_sub_squares(self):
-        mat = [[1, 1, 1, 1, 1], 
-           [2, 2, 2, 2, 2], 
-           [3, 3, 3, 3, 3], 
-           [4, 4, 4, 4, 4], 
-           [5, 5, 5, 5, 5]]
+        mat = [[1, 1, 1, 1, 1],
+               [2, 2, 2, 2, 2],
+               [3, 3, 3, 3, 3],
+               [4, 4, 4, 4, 4],
+               [5, 5, 5, 5, 5]]
         self.assertEqual(sum_sub_squares.sum_sub_squares(mat, 3),
                          [[18, 18, 18], [27, 27, 27], [36, 36, 36]])
 
+
+class TestSortMatrixDiagonally(unittest.TestCase):
+    def test_sort_diagonally(self):
+        mat = [
+            [3, 3, 1, 1],
+            [2, 2, 1, 2],
+            [1, 1, 1, 2]
+        ]
+
+        self.assertEqual(sort_matrix_diagonally.sort_diagonally(mat), [
+            [1, 1, 1, 1],
+            [1, 2, 2, 2],
+            [1, 2, 3, 3]
+        ])
+
+
 if __name__ == "__main__":
     unittest.main()