Add solution for challenge #58

Author: Mauricio Klein

README.md

@@ -57,3 +57,4 @@
 - [x] [Challenge 55: Merge list of number into ranges](challenge-55/)
 - [x] [Challenge 56: Find largest BST](challenge-56/)
 - [x] [Challenge 57: Arithmetic Binary Tree](challenge-57/)
+- [x] [Challenge 58: Get all Values at a Certain Height in a Binary Tree](challenge-58/)

challenge-58/Makefile

@@ -0,0 +1,4 @@
+test:
+	python test_*.py
+
+.PHONY: test

challenge-58/README.md

@@ -0,0 +1,22 @@
+# Get all Values at a Certain Height in a Binary Tree
+
+This problem was asked by Amazon.
+
+## Description
+
+Given a binary tree, return all values given a certain height h.
+
+## Example
+
+```
+    1
+   / \
+  2   3
+ / \   \
+4   5   7
+
+h = 3
+
+Output: [4, 5, 7]
+```
+

challenge-58/__init__.py

@@ -0,0 +1,21 @@
+class Node(object):
+  def __init__(self, v, left=None, right=None):
+    self.v = v
+    self.left = left
+    self.right = right
+
+#
+# Time  complexity: O(n) (every element is visited at most once)
+# Space complexity: O(n) (worst case, considering "height" equal to the tree's depth)
+#
+def values_at_height(node, height):
+  if not node:
+    return []
+
+  if height == 1:
+    return [node.v]
+  else:
+    left  = values_at_height(node.left,  height - 1)
+    right = values_at_height(node.right, height - 1)
+
+    return left + right

challenge-58/solver.py

@@ -0,0 +1,23 @@
+import unittest
+import solver
+from solver import Node, values_at_height
+
+class TestSolver(unittest.TestCase):
+  def test_values_at_height(self):
+    root = Node(1)
+    self.assertEqual(values_at_height(root, 1), [1])
+
+    root = Node(1, Node(2), Node(3))
+    self.assertEqual(values_at_height(root, 2), [2, 3])
+    
+    root = Node(1, Node(3))
+    self.assertEqual(values_at_height(root, 2), [3])
+
+    root = Node(1, Node(2), Node(3))
+    self.assertEqual(values_at_height(root, 5), [])
+
+    root = Node(1, Node(2, Node(4), Node(5)), Node(3, right=Node(7)))
+    self.assertEqual(values_at_height(root, 3), [4, 5, 7])
+    
+if __name__ == "__main__":
+    unittest.main()