enhanced segment tree implementation and more pythonic

data_structures/binary_tree/segment_tree_other.py

@@ -138,9 +138,34 @@ def __init__(self, collection: Sequence, function):
             self.root = self._build_tree(0, len(collection) - 1)
 
     def update(self, i, val):
+        """
+        update value in collection
+        :param i: index of collection
+        :param val: new value
+        :return:
+        >>> import operator
+        >>> num_arr = SegmentTree([2, 1, 5, 3, 4], operator.add)
+        >>> num_arr.update(1, 5)
+        """
         self._update_tree(self.root, i, val)
 
     def query_range(self, i, j):
+        """
+        Sum, Max, Min operation in intervals i and j([i, j])
+        :param i:  left index
+        :param j:  right index
+        :return:  Sum, Max, Min
+        >>> import operator
+        >>> num_arr = SegmentTree([2, 1, 5, 3, 4], operator.add)
+        >>> num_arr.update(1, 5)
+        >>> num_arr.query_range(3, 4)
+        7
+        >>> num_arr.query_range(2, 2)
+        5
+        >>> num_arr.query_range(1, 3)
+        13
+        >>>
+        """
         return self._query_range(self.root, i, j)
 
     def _build_tree(self, start, end):
@@ -176,20 +201,18 @@ def _query_range(self, root, i, j):
             return self.fn(self._query_range(root.left, i, root.mid), self._query_range(root.right, root.mid + 1, j))
 
     def traverse(self):
-        result = []
         if self.root is not None:
             queue = Queue()
             queue.put(self.root)
             while not queue.empty():
                 node = queue.get()
-                result.append(node)
+                yield node
 
                 if node.left is not None:
                     queue.put(node.left)
 
                 if node.right is not None:
                     queue.put(node.right)
-            return result
 
 
 if __name__ == '__main__':