Add method 'bisect_right' and test

fenwick/fenwick.py

@@ -93,6 +93,18 @@ def init(self, frequencies):
             parent_idx = idx + (idx & -idx)  # parent in update tree
             if parent_idx <= self._n:
                 self._v[parent_idx - 1] += self._v[idx - 1]
+
+    def bisect_right(self, value):
+        """Returns the smallest index i such that the cumulative sum up to i is > value."""
+        # https://stackoverflow.com/questions/34699616/fenwick-trees-to-determine-which-interval-a-point-falls-in
+        j = 2 ** len(self)
+        idx = -1
+        while j > 0:
+            if idx + j < len(self) and value >= self._v[idx + j]:
+                value -= self._v[idx + j]
+                idx += j
+            j >>= 1
+        return idx + 1
 
     def __eq__(self, other):
         return isinstance(other, FenwickTree) and self._n == other._n and self._v == other._v

tests/test_fenwick.py

@@ -1,4 +1,6 @@
 import unittest
+import itertools
+import bisect
 
 from fenwick import FenwickTree
 
@@ -34,6 +36,29 @@ def test_fenwick(self):
             self.assertEqual(fenwick_tree_1[idx], frequencies[idx])
 
         self.assertEqual(fenwick_tree_0, fenwick_tree_1)
+
+
+    def test_bisect_left(self):
+
+        # Create a Fenwick tree
+        frequencies = [1, 2, 3]
+        n = len(frequencies)
+        tree = FenwickTree(n)
+        tree.init(frequencies)
+
+        # Create a cumulative sum
+        cumsum = list(itertools.accumulate(frequencies))
+
+        # Test at some random values, including edge cases
+        test_values = [-5, 0.5, 1, 2, 3, 7]
+        for test_value in test_values:
+            index_fenwick = tree.bisect_right(test_value)
+            index_bisect = bisect.bisect_right(cumsum, test_value)
+            self.assertEqual(index_fenwick, index_bisect)
+
+        # Test a specific case. The cumsums are [1, 3, 6], and the 
+        # smallest index i such that cumsum[i] > 3 is 2.
+        self.assertEqual(tree.bisect_right(3), 2)
 
 
 if __name__ == '__main__':