Merge pull request #1 from leetcode/master

1e0ng · web-flow · commit ad2840d572a1 · 2019-02-15T12:01:14.000+08:00
Add out of bounds check for segment tree
diff --git a/segmenttree/__init__.py b/segmenttree/__init__.py
@@ -8,19 +8,36 @@ def __init__(self, start, end):
         self.len_value = {}
         self._init(start, end)
 
-    def add(self, start, end, weight=1):
+    def out_of_bounds_check(self, start, end):
         start = max(start, self.start)
         end = min(end, self.end)
+        if start > end:
+            return None, None
+        return start, end
+
+    def add(self, start, end, weight=1):
+        start, end = self.out_of_bounds_check(start, end)
+        if start is None:
+            return False
         self._add(start, end, weight, self.start, self.end)
         return True
 
     def query_max(self, start, end):
+        start, end = self.out_of_bounds_check(start, end)
+        if start is None:
+            return None
         return self._query_max(start, end, self.start, self.end)
 
     def query_sum(self, start, end):
+        start, end = self.out_of_bounds_check(start, end)
+        if start is None:
+            return 0
         return self._query_sum(start, end, self.start, self.end)
 
     def query_len(self, start, end):
+        start, end = self.out_of_bounds_check(start, end)
+        if start is None:
+            return 0
         return self._query_len(start, end, self.start, self.end)
 
     """"""
diff --git a/tests/test_segtree.py b/tests/test_segtree.py
@@ -50,5 +50,45 @@ def test_empty(self):
         self.assertEqual(2, segtree.query_len(0, 8))
         self.assertEqual(2, segtree.query_sum(0, 8))
 
+    def test_full_out_of_bound(self):
+        segtree = SegmentTree(0, 8)
+        segtree.add(0, 8)
+
+        # Test full out of bound adding element fails
+        self.assertEqual(False, segtree.add(10, 16))
+        self.assertEqual(False, segtree.add(-16, -10))
+
+        # Test full out of bound len query returns 0
+        self.assertEqual(0, segtree.query_len(10, 16))
+        self.assertEqual(0, segtree.query_len(-16, -10))
+
+        # Test full out of bounds len query returns 0
+        self.assertEqual(0, segtree.query_sum(10, 16))
+        self.assertEqual(0, segtree.query_sum(-16, -10))
+
+        # Test full out of bounds max query returns None
+        self.assertEqual(None, segtree.query_max(10, 16))
+        self.assertEqual(None, segtree.query_max(-16, -10))
+
+    def test_partial_out_of_bound(self):
+        segtree = SegmentTree(0, 8)
+        segtree.add(0, 8)
+
+        # Test partial out of bound adding element
+        self.assertEqual(True, segtree.add(8, 16))
+        self.assertEqual(True, segtree.add(-16, 0))
+
+        # Test partial out of bound len query
+        self.assertEqual(2, segtree.query_len(7, 16))
+        self.assertEqual(2, segtree.query_len(-16, 1))
+
+        # Test partial out of bounds len query
+        self.assertEqual(3, segtree.query_sum(7, 16))
+        self.assertEqual(3, segtree.query_sum(-16, 1))
+
+        # Test full out of bounds max query
+        self.assertEqual(2, segtree.query_max(7, 16))
+        self.assertEqual(2, segtree.query_max(-16, 1))
+
 if __name__ == '__main__':
     main()
