LC 1372. Longest ZigZag Path in a Binary Tree (Python DFS)

Author: raul-sauco

README.md

@@ -394,6 +394,7 @@ Solutions to LeetCode problems. The first column links to the problem in LeetCod
 | [1343. Number of Sub-arrays of Size K and Average Greater than or Equal to Threshold][lc1343] | 🟠 Medium  | [![python](res/py.png)][lc1343py] [![rust](res/rs.png)][lc1343rs]                                     |
 | [1345. Jump Game IV][lc1345]                                                                  | 🔴 Hard    | [![python](res/py.png)][lc1345py] [![rust](res/rs.png)][lc1345rs]                                     |
 | [1354. Construct Target Array With Multiple Sums][lc1354]                                     | 🔴 Hard    | [![python](res/py.png)][lc1354py]                                                                     |
+| [1372. Longest ZigZag Path in a Binary Tree][lc1372]                                          | 🟠 Medium  | [![python](res/py.png)][lc1372py]                                                                     |
 | [1383. Maximum Performance of a Team][lc1383]                                                 | 🔴 Hard    | [![python](res/py.png)][lc1383py]                                                                     |
 | [1402. Reducing Dishes][lc1402]                                                               | 🔴 Hard    | [![python](res/py.png)][lc1402py] [![rust](res/rs.png)][lc1402rs]                                     |
 | [1423. Maximum Points You Can Obtain from Cards][lc1423]                                      | 🟠 Medium  | [![python](res/py.png)][lc1423py]                                                                     |
@@ -1310,6 +1311,8 @@ Solutions to LeetCode problems. The first column links to the problem in LeetCod
 [lc1345rs]: leetcode/jump-game-iv.rs
 [lc1354]: https://leetcode.com/problems/construct-target-array-with-multiple-sums/
 [lc1354py]: leetcode/construct-target-array-with-multiple-sums.py
+[lc1372]: https://leetcode.com/problems/longest-zigzag-path-in-a-binary-tree/
+[lc1372py]: leetcode/longest-zigzag-path-in-a-binary-tree.py
 [lc1383]: https://leetcode.com/problems/maximum-performance-of-a-team/
 [lc1383py]: leetcode/maximum-performance-of-a-team.py
 [lc1402]: https://leetcode.com/problems/reducing-dishes/

leetcode/longest-zigzag-path-in-a-binary-tree.py

@@ -0,0 +1,113 @@
+# 1372. Longest ZigZag Path in a Binary Tree
+# 🟠 Medium
+#
+# https://leetcode.com/problems/longest-zigzag-path-in-a-binary-tree/
+#
+# Tags: Dynamic Programming - Tree - Depth-First Search - Binary Tree
+
+import timeit
+from enum import Enum
+from typing import Optional
+
+from utils.binary_tree import BinaryTree, TreeNode
+
+
+# We could use an enum to make the code easier to read but it makes it
+# a little less efficient because it need to create the object.
+class Direction(Enum):
+    RIGHT = "right"
+    LEFT = "left"
+
+
+# Use any traversal method to visit all nodes, for each node, record the
+# direction that we used to travel there and the length of the zig-zag
+# path to it, the path will continue to grow for one of its children and
+# will restart for the other one.
+#
+# Time complexity: O(n) - We visit all nodes and do O(1) work for each.
+# Space complexity: O(n) - The stack can grow to size n.
+#
+# Runtime 367 ms Beats 89.1%
+# Memory 26.7 MB Beats 95.60%
+class Solution:
+    def longestZigZag(self, root: Optional[TreeNode]) -> int:
+        if not root:
+            return 0
+        # Traverse the tree, for each node, keep its position in a zig
+        # zag path and the direction that we traveled to get there.
+        stack, res = [], 0
+        if root.left:
+            stack.append((root.left, Direction.LEFT, 1))
+        if root.right:
+            stack.append((root.right, Direction.RIGHT, 1))
+        while stack:
+            current, last_dir, path_length = stack.pop()
+            if path_length > res:
+                res = path_length
+            if current.left:
+                stack.append(
+                    (
+                        current.left,
+                        Direction.LEFT,
+                        path_length + 1 if last_dir == Direction.RIGHT else 1,
+                    )
+                )
+            if current.right:
+                stack.append(
+                    (
+                        current.right,
+                        Direction.RIGHT,
+                        path_length + 1 if last_dir == Direction.LEFT else 1,
+                    )
+                )
+        return res
+
+
+def test():
+    executors = [Solution]
+    tests = [
+        [[1], 0],
+        [[1, 1, 1, None, 1, None, None, 1, 1, None, 1], 4],
+        [
+            [
+                1,
+                None,
+                1,
+                1,
+                1,
+                None,
+                None,
+                1,
+                1,
+                None,
+                1,
+                None,
+                None,
+                None,
+                1,
+                None,
+                1,
+            ],
+            3,
+        ],
+    ]
+    for executor in executors:
+        start = timeit.default_timer()
+        for _ in range(1):
+            for col, t in enumerate(tests):
+                sol = executor()
+                root = BinaryTree.fromList(t[0]).getRoot()
+                result = sol.longestZigZag(root)
+                exp = t[1]
+                assert result == exp, (
+                    f"\033[93m» {result} <> {exp}\033[91m for"
+                    + f" test {col} using \033[1m{executor.__name__}"
+                )
+        stop = timeit.default_timer()
+        used = str(round(stop - start, 5))
+        cols = "{0:20}{1:10}{2:10}"
+        res = cols.format(executor.__name__, used, "seconds")
+        print(f"\033[92m» {res}\033[0m")
+
+
+test()