Skip to content

feat: Implement Fenwich Tree (Binary Indexed Tree) and Implement Additional Trie Operations #649

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 12 commits into from
1 change: 1 addition & 0 deletions AUTHORS
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,4 @@ Pratik Goyal <[email protected]>
Jay Thorat <[email protected]>
Rajveer Singh Bharadwaj <[email protected]>
Kishan Ved <[email protected]>
Arvinder Singh Dhoul <[email protected]>
89 changes: 89 additions & 0 deletions pydatastructs/trees/fenwich_tree.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
__all__ = [
'fenwich_tree'
]

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)
112 changes: 112 additions & 0 deletions pydatastructs/trees/tests/test_fenwich_tree.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
import unittest
from pydatastructs.trees.fenwich_tree 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
111 changes: 111 additions & 0 deletions pydatastructs/trees/tests/test_trie.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
import pytest
from pydatastructs.trees.trie import Trie

def test_trie_insert_search():
trie = Trie()
trie.insert("apple")
assert trie.search("apple")
assert not trie.search("app")
trie.insert("app")
assert trie.search("app")

def test_trie_starts_with():
trie = Trie()
trie.insert("apple")
assert trie.starts_with("app")
assert trie.starts_with("a")
assert not trie.starts_with("b")
assert not trie.starts_with("applxyz")

def test_trie_empty():
trie = Trie()
assert not trie.search("apple")
assert not trie.starts_with("app")

def test_trie_multiple_words():
trie = Trie()
trie.insert("apple")
trie.insert("application")
trie.insert("banana")
assert trie.search("apple")
assert trie.search("application")
assert trie.search("banana")
assert not trie.search("app")
assert trie.starts_with("app")
assert trie.starts_with("ban")
assert not trie.starts_with("aplx")

def test_trie_case_sensitive():
trie = Trie()
trie.insert("Apple")
assert trie.search("Apple")
assert not trie.search("apple")

def test_count_words():
trie = Trie()
assert trie.count_words() == 0
trie.insert("apple")
assert trie.count_words() == 1
trie.insert("app")
assert trie.count_words() == 2
trie.insert("apple")
assert trie.count_words() == 2

def test_longest_common_prefix():
trie = Trie()
assert trie.longest_common_prefix() == ""
trie.insert("apple")
assert trie.longest_common_prefix() == "apple"
trie.insert("application")
assert trie.longest_common_prefix() == "appl"
trie.insert("banana")
assert trie.longest_common_prefix() == ""

def test_autocomplete():
trie = Trie()
trie.insert("apple")
trie.insert("application")
trie.insert("app")
assert trie.autocomplete("app") == ["app", "apple", "application"]
assert trie.autocomplete("appl") == ["apple", "application"]
assert trie.autocomplete("b") == []

def test_bulk_insert():
trie = Trie()
trie.bulk_insert(["apple", "banana", "orange"])
assert trie.search("apple")
assert trie.search("banana")
assert trie.search("orange")
assert trie.count_words() == 3

def test_clear():
trie = Trie()
trie.insert("apple")
trie.clear()
assert trie.is_empty()
assert trie.count_words() == 0
assert not trie.search("apple")

def test_is_empty():
trie = Trie()
assert trie.is_empty()
trie.insert("apple")
assert not trie.is_empty()
trie.clear()
assert trie.is_empty()

def test_find_all_words():
trie = Trie()
trie.bulk_insert(["apple", "banana", "orange"])
assert sorted(trie.find_all_words()) == sorted(["apple", "banana", "orange"])
trie.clear()
assert trie.find_all_words() == []


def test_longest_word():
trie = Trie()
assert trie.longest_word() is None
trie.bulk_insert(["apple", "banana", "application"])
assert trie.longest_word() == "application"
trie.insert("a")
assert trie.longest_word() == "application"
Loading
Loading