[codex] Add optimized diagonal traverse solution

practice/leetcode/solutions_00300/solution_00498.py

@@ -35,12 +35,77 @@ def findDiagonalOrder(self, mat):
         return res
 
 
+class Solution2(object):
+    def findDiagonalOrder(self, mat):
+        """
+        :type mat: List[List[int]]
+        :rtype: List[int]
+        """
+        rows = len(mat)
+        cols = len(mat[0])
+        result = []
+
+        row = 0
+        col = 0
+        direction = 1
+
+        for _ in range(rows * cols):
+            result.append(mat[row][col])
+
+            if direction == 1:
+                if col == cols - 1:
+                    row += 1
+                    direction = -1
+                elif row == 0:
+                    col += 1
+                    direction = -1
+                else:
+                    row -= 1
+                    col += 1
+            else:
+                if row == rows - 1:
+                    col += 1
+                    direction = 1
+                elif col == 0:
+                    row += 1
+                    direction = 1
+                else:
+                    row += 1
+                    col -= 1
+
+        return result
+
+
 class TestSolution(unittest.TestCase):
     def test_findDiagonalOrder(self):
-        solution = Solution()
+        cases = [
+            (
+                [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                [1, 2, 4, 7, 5, 3, 6, 8, 9],
+            ),
+            (
+                [[1, 2], [3, 4]],
+                [1, 2, 3, 4],
+            ),
+            (
+                [[1, 2, 3], [4, 5, 6]],
+                [1, 2, 4, 5, 3, 6],
+            ),
+            (
+                [[1, 2], [3, 4], [5, 6]],
+                [1, 2, 3, 5, 4, 6],
+            ),
+            (
+                [[1]],
+                [1],
+            ),
+        ]
+
+        for solution_class in (Solution, Solution2):
+            solution = solution_class()
 
-        self.assertListEqual(solution.findDiagonalOrder([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
-                             [1, 2, 4, 7, 5, 3, 6, 8, 9])
+            for mat, expected in cases:
+                self.assertListEqual(solution.findDiagonalOrder(mat), expected)
 
 
 if __name__ == '__main__':