Added Segment Trees (#10)

Author: czgdp1807

.travis.yml

@@ -6,4 +6,3 @@ install:
         - pip install -r requirements.txt
 script:
         - pytest --doctest-modules
-        - pytest

pydatastructs/trees/__init__.py

@@ -5,3 +5,9 @@
     Node, BinaryTree, BinarySearchTree
 )
 __all__.extend(binary_trees.__all__)
+
+from . import space_partitioning_trees
+from .space_partitioning_trees import (
+    OneDimensionalSegmentTree
+)
+__all__.extend(space_partitioning_trees.__all__)

pydatastructs/trees/space_partitioning_trees.py

@@ -0,0 +1,229 @@
+from pydatastructs.trees.binary_trees import Node
+from collections import deque as Queue
+from pydatastructs.linear_data_structures.arrays import _check_type
+
+__all__ = [
+    'OneDimensionalSegmentTree'
+]
+
+class OneDimensionalSegmentTree(object):
+    """
+    Represents one dimensional segment trees.
+
+    Parameters
+    ==========
+
+    segs: list/tuple/set
+        The segs should contains tuples/list/set of size 2
+        denoting the start and end points of the intervals.
+
+    Examples
+    ========
+
+    >>> from pydatastructs import OneDimensionalSegmentTree as ODST
+    >>> segt = ODST([(3, 8), (9, 20)])
+    >>> segt.build()
+    >>> segt.tree[0].key
+    [False, 2, 3, False]
+    >>> len(segt.query(4))
+    1
+
+    Note
+    ====
+
+    All the segments are assumed to be closed intervals,
+    i.e., the ends are points of segments are also included in
+    computation.
+
+    References
+    ==========
+
+    .. [1] https://en.wikipedia.org/wiki/Segment_tree
+
+    """
+
+    __slots__ = ['segments', 'tree', 'root_idx', 'cache']
+
+    def __new__(cls, segs):
+        obj = object.__new__(cls)
+        if any((not isinstance(seg, (tuple, list, set)) or len(seg) != 2)
+                for seg in segs):
+                    raise ValueError('%s is invalid set of intervals'%(segs))
+        for i in range(len(segs)):
+            segs[i] = list(segs[i])
+            segs[i].sort()
+        obj.segments = [seg for seg in segs]
+        obj.tree, obj.root_idx, obj.cache = [], None, False
+        return obj
+
+    def _union(self, i1, i2):
+        """
+        Helper function for taking union of two
+        intervals.
+        """
+        return Node([i1.key[0], i1.key[1], i2.key[2], i2.key[3]], None)
+
+    def _intersect(self, i1, i2):
+        """
+        Helper function for finding intersection of two
+        intervals.
+        """
+        if i1 == None or i2 == None:
+            return False
+        if i1.key[2] < i2.key[1] or i2.key[2] < i1.key[1]:
+            return False
+        c1, c2 = None, None
+        if i1.key[2] == i2.key[1]:
+            c1 = (i1.key[3] and i2.key[0])
+        if i2.key[2] == i1.key[1]:
+            c2 = (i2.key[3] and i1.key[0])
+        if c1 == False and c2 == False:
+            return False
+        return True
+
+    def _contains(self, i1, i2):
+        """
+        Helper function for checking if the first interval
+        is contained in second interval.
+        """
+        if i1 == None or i2 == None:
+            return False
+        if i1.key[1] < i2.key[1] and i1.key[2] > i2.key[2]:
+            return True
+        if i1.key[1] == i2.key[1] and i1.key[2] > i2.key[2]:
+            return (i1.key[0] or not i2.key[0])
+        if i1.key[1] < i2.key[1] and i1.key[2] == i2.key[2]:
+            return i1.key[3] or not i2.key[3]
+        if i1.key[1] == i2.key[1] and i1.key[2] == i2.key[2]:
+            return not ((not i1.key[3] and i2.key[3]) or (not i1.key[0] and i2.key[0]))
+        return False
+
+    def _iterate(self, calls, I, idx):
+        """
+        Helper function for filling the calls
+        stack. Used for imitating the stack based
+        approach used in recursion.
+        """
+        if self.tree[idx].right == None:
+            rc = None
+        else:
+            rc = self.tree[self.tree[idx].right]
+        if self.tree[idx].left == None:
+            lc = None
+        else:
+            lc = self.tree[self.tree[idx].left]
+        if self._intersect(I, rc):
+            calls.append(self.tree[idx].right)
+        if self._intersect(I, lc):
+            calls.append(self.tree[idx].left)
+        return calls
+
+    def build(self):
+        """
+        Builds the segment tree from the segments,
+        using iterative algorithm based on stacks.
+        """
+        if self.cache:
+            return None
+        endpoints = []
+        for segment in self.segments:
+            endpoints.extend(segment)
+        endpoints.sort()
+
+        elem_int = Queue()
+        elem_int.append(Node([False, endpoints[0] - 1, endpoints[0], False], None))
+        i = 0
+        while i < len(endpoints) - 1:
+            elem_int.append(Node([True, endpoints[i], endpoints[i], True], None))
+            elem_int.append(Node([False, endpoints[i], endpoints[i+1], False], None))
+            i += 1
+        elem_int.append(Node([True, endpoints[i], endpoints[i], True], None))
+        elem_int.append(Node([False, endpoints[i], endpoints[i] + 1, False], None))
+
+        self.tree = []
+        while len(elem_int) > 1:
+            m = len(elem_int)
+            while m >= 2:
+                I1 = elem_int.popleft()
+                I2 = elem_int.popleft()
+                I = self._union(I1, I2)
+                I.left = len(self.tree)
+                I.right = len(self.tree) + 1
+                self.tree.append(I1), self.tree.append(I2)
+                elem_int.append(I)
+                m -= 2
+            if m & 1 == 1:
+                Il = elem_int.popleft()
+                elem_int.append(Il)
+
+        Ir = elem_int.popleft()
+        Ir.left, Ir.right = -3, -2
+        self.tree.append(Ir)
+        self.root_idx = -1
+
+        for segment in self.segments:
+            I = Node([True, segment[0], segment[1], True], None)
+            calls = [self.root_idx]
+            while calls:
+                idx = calls.pop()
+                if self._contains(I, self.tree[idx]):
+                    if self.tree[idx].data == None:
+                        self.tree[idx].data = []
+                    self.tree[idx].data.append(I)
+                    continue
+                calls = self._iterate(calls, I, idx)
+        self.cache = True
+
+    def query(self, qx, init_node=None):
+        """
+        Queries the segment tree.
+
+        Parameters
+        ==========
+
+        qx: int/float
+            The query point
+        init_node: int
+            The index of the node from which the query process
+            is to be started.
+
+        Returns
+        =======
+
+        intervals: set
+            The set of the intervals which contain the query
+            point.
+
+        References
+        ==========
+
+        .. [1] https://en.wikipedia.org/wiki/Segment_tree
+        """
+        if not self.cache:
+            self.build()
+        if init_node == None:
+            init_node = self.root_idx
+        qn = Node([True, qx, qx, True], None)
+        intervals = []
+        calls = [init_node]
+        while calls:
+            idx = calls.pop()
+            if _check_type(self.tree[idx].data, list):
+                intervals.extend(self.tree[idx].data)
+            calls = self._iterate(calls, qn, idx)
+        return set(intervals)
+
+    def __str__(self):
+        """
+        Used for printing.
+        """
+        if not self.cache:
+            self.build()
+        str_tree = []
+        for seg in self.tree:
+            if seg.data == None:
+                data = None
+            else:
+                data = [str(sd) for sd in seg.data]
+            str_tree.append((seg.left, seg.key, data, seg.right))
+        return str(str_tree)

pydatastructs/trees/test_space_partitioning_tree.py

@@ -0,0 +1,21 @@
+from pydatastructs import OneDimensionalSegmentTree
+from pydatastructs.utils.raises_util import raises
+
+def test_OneDimensionalSegmentTree():
+    ODST = OneDimensionalSegmentTree
+    segt = ODST([(0, 5), (1, 6), (9, 13), (1, 2), (3, 8), (9, 20)])
+    assert segt.cache == False
+    segt.build()
+    assert segt.cache == True
+    segt2 = ODST([(1, 4)])
+    assert str(segt2) == ("[(None, [False, 0, 1, False], None, None), "
+    "(None, [True, 1, 1, True], ['(None, [True, 1, 4, True], None, None)'], "
+    "None), (None, [False, 1, 4, False], None, None), (None, [True, 4, 4, True], "
+    "None, None), (0, [False, 0, 1, True], None, 1), (2, [False, 1, 4, True], "
+    "['(None, [True, 1, 4, True], None, None)'], 3), (4, [False, 0, 4, True], "
+    "None, 5), (None, [False, 4, 5, False], None, None), (-3, [False, 0, 5, "
+    "False], None, -2)]")
+    assert len(segt.query(1.5)) == 3
+    assert len(segt.query(-1)) == 0
+    assert len(segt.query(2.8)) == 2
+    raises(ValueError, lambda: ODST([(1, 2, 3)]))