diff --git a/practice/leetcode/solutions_00300/solution_00498.py b/practice/leetcode/solutions_00300/solution_00498.py index 9669d7b..ef8cfa5 100644 --- a/practice/leetcode/solutions_00300/solution_00498.py +++ b/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__':