added feature: fenwich tree

Author: arvinder004

pydatastructs/trees/fenwich_tree.py

@@ -0,0 +1,85 @@
+class fenwich_tree:
+    """
+    Implementation of Fenwich tree/Binary Indexed Tree
+    """
+
+    def __init__(self, size_or_array):
+        """
+        Initializes the Fenwich Tree.
+
+        Args:
+            size_or_array: size of array the tree will represent or array of values
+        """
+
+        if isinstance(size_or_array, int):
+            self.size = size_or_array
+            self.tree = [0] * (self.size + 1)
+            self.original_array = [0] * self.size 
+        elif isinstance(size_or_array, list):
+            self.original_array = list(size_or_array)
+            self.size = len(self.original_array)
+            self.tree = [0] * (self.size + 1)
+            for i, val in enumerate(self.original_array):
+                self._update_tree(i, val)
+        else:
+            raise ValueError("size_or_array must be an integer or a list.")
+        
+    def _update_tree(self, index, delta):
+        """
+        Internal helper to update the Fenwick Tree after a change in the original array.
+        """
+        index += 1  # Fenwick Tree is 1-indexed
+        while index <= self.size:
+            self.tree[index] += delta
+            index += index & (-index)
+
+    def update(self, index, value):
+        """
+        Updates the value at the given index in the original array and the Fenwick Tree.
+
+        Args:
+            index: The index to update (0-based).
+            value: The new value.
+        """
+        if not (0 <= index < self.size):
+            raise IndexError("Index out of bounds")
+        delta = value - self.original_array[index]
+        self.original_array[index] = value
+        self._update_tree(index, delta)
+
+    def prefix_sum(self, index):
+        """
+        Calculates the prefix sum up to the given index (inclusive).
+
+        Args:
+            index: The index up to which to calculate the sum (0-based).
+
+        Returns:
+            The prefix sum.
+        """
+        if not (0 <= index < self.size):
+            raise IndexError("Index out of bounds")
+        index += 1  #
+        sum_val = 0
+        while index > 0:
+            sum_val += self.tree[index]
+            index -= index & (-index)  
+        return sum_val
+    
+    def range_sum(self, start_index, end_index):
+        """
+        Calculates the sum of elements within the given range (inclusive).
+
+        Args:
+            start_index: The starting index of the range (0-based).
+            end_index: The ending index of the range (0-based).
+
+        Returns:
+            The sum of elements in the range.
+        """
+        if not (0 <= start_index <= end_index < self.size):
+            raise IndexError("Indices out of bounds")
+        if start_index == 0:
+            return self.prefix_sum(end_index)
+        else:
+            return self.prefix_sum(end_index) - self.prefix_sum(start_index - 1)

pydatastructs/trees/tests/test_fenwich_tree.py

@@ -0,0 +1,112 @@
+import unittest
+from pydatastructs import fenwich_tree
+
+class TestFenwickTree(unittest.TestCase):
+
+    def test_initialization_with_size(self):
+        ft = fenwich_tree(5)
+        self.assertEqual(ft.size, 5)
+        self.assertEqual(ft.tree, [0, 0, 0, 0, 0, 0])
+        self.assertEqual(ft.original_array, [0, 0, 0, 0, 0])
+
+    def test_initialization_with_array(self):
+        arr = [1, 2, 3, 4, 5]
+        ft = fenwich_tree(arr)
+        self.assertEqual(ft.size, 5)
+        self.assertEqual(ft.original_array, arr)
+        # Manually calculate prefix sums and check the tree structure
+        expected_tree = [0, 1, 3, 3, 10, 5]
+        self.assertEqual(ft.tree, expected_tree)
+
+    def test_initialization_with_empty_array(self):
+        arr = []
+        ft = fenwich_tree(arr)
+        self.assertEqual(ft.size, 0)
+        self.assertEqual(ft.tree, [0])
+        self.assertEqual(ft.original_array, [])
+
+    def test_initialization_with_invalid_input(self):
+        with self.assertRaises(ValueError):
+            fenwich_tree("invalid")
+
+    def test_update_single_element(self):
+        ft = fenwich_tree([1, 2, 3, 4, 5])
+        ft.update(1, 10)
+        self.assertEqual(ft.original_array, [1, 10, 3, 4, 5])
+        expected_tree = [0, 1, 11, 3, 18, 5]
+        self.assertEqual(ft.tree, expected_tree)
+
+    def test_update_out_of_bounds(self):
+        ft = fenwich_tree(5)
+        with self.assertRaises(IndexError):
+            ft.update(5, 10)
+        with self.assertRaises(IndexError):
+            ft.update(-1, 10)
+
+    def test_prefix_sum_positive_indices(self):
+        arr = [1, 2, 3, 4, 5]
+        ft = fenwich_tree(arr)
+        self.assertEqual(ft.prefix_sum(0), 1)
+        self.assertEqual(ft.prefix_sum(1), 3)
+        self.assertEqual(ft.prefix_sum(2), 6)
+        self.assertEqual(ft.prefix_sum(3), 10)
+        self.assertEqual(ft.prefix_sum(4), 15)
+
+    def test_prefix_sum_out_of_bounds(self):
+        ft = fenwich_tree(5)
+        with self.assertRaises(IndexError):
+            ft.prefix_sum(5)
+        with self.assertRaises(IndexError):
+            ft.prefix_sum(-1)
+
+    def test_prefix_sum_empty_array(self):
+        ft = fenwich_tree([])
+        with self.assertRaises(IndexError):
+            ft.prefix_sum(0) # Should raise IndexError as size is 0
+
+    def test_range_sum_valid_range(self):
+        arr = [1, 2, 3, 4, 5]
+        ft = fenwich_tree(arr)
+        self.assertEqual(ft.range_sum(0, 0), 1)
+        self.assertEqual(ft.range_sum(0, 1), 3)
+        self.assertEqual(ft.range_sum(1, 3), 2 + 3 + 4)
+        self.assertEqual(ft.range_sum(2, 4), 3 + 4 + 5)
+        self.assertEqual(ft.range_sum(0, 4), 1 + 2 + 3 + 4 + 5)
+
+    def test_range_sum_out_of_bounds(self):
+        ft = fenwich_tree(5)
+        with self.assertRaises(IndexError):
+            ft.range_sum(0, 5)
+        with self.assertRaises(IndexError):
+            ft.range_sum(-1, 2)
+        with self.assertRaises(IndexError):
+            ft.range_sum(1, 5)
+        with self.assertRaises(IndexError):
+            ft.range_sum(-1, -1)
+
+    def test_range_sum_invalid_range(self):
+        ft = fenwich_tree(5)
+        with self.assertRaises(IndexError):
+            ft.range_sum(3, 1)
+
+    def test_range_sum_single_element(self):
+        arr = [10, 20, 30]
+        ft = fenwich_tree(arr)
+        self.assertEqual(ft.range_sum(0, 0), 10)
+        self.assertEqual(ft.range_sum(1, 1), 20)
+        self.assertEqual(ft.range_sum(2, 2), 30)
+
+    def test_range_sum_entire_array(self):
+        arr = [1, 2, 3, 4, 5]
+        ft = fenwich_tree(arr)
+        self.assertEqual(ft.range_sum(0, ft.size - 1), 15)
+
+    def test_update_and_query_sequence(self):
+        ft = fenwich_tree([2, 5, 1, 8, 3])
+        self.assertEqual(ft.prefix_sum(3), 2 + 5 + 1 + 8)  # 16
+        ft.update(1, 10)
+        self.assertEqual(ft.prefix_sum(3), 2 + 10 + 1 + 8) # 21
+        self.assertEqual(ft.range_sum(0, 2), 2 + 10 + 1)   # 13
+        ft.update(4, 0)
+        self.assertEqual(ft.prefix_sum(4), 2 + 10 + 1 + 8 + 0) # 21
+        self.assertEqual(ft.range_sum(3, 4), 8 + 0)       # 8