Skip to content

Commit 57ffc71

Browse files
committed
feat: fenwick tree
Add Fenwick tree implmentation and tests
1 parent 9545dd3 commit 57ffc71

File tree

5 files changed

+234
-15
lines changed

5 files changed

+234
-15
lines changed

README.md

Lines changed: 3 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,10 @@
33
[![Build](https://github.com/rubenperezm/pystrukts/actions/workflows/test.yml/badge.svg)](https://github.com/rubenperezm/pystrukts/actions/workflows/test.yml/badge.svg)
44
[![codecov](https://codecov.io/gh/rubenperezm/pystrukts/graph/badge.svg?token=OLV3EOPYFI)](https://codecov.io/gh/rubenperezm/pystrukts)
55
[![PyPI Version](https://img.shields.io/pypi/v/pystrukts.svg)](https://pypi.org/project/pystrukts)
6-
<!-- [![Python Versions](https://img.shields.io/pypi/pyversions/pystrukts.svg)](https://pypi.org/project/pystrukts) -->
6+
[![Python Versions](https://img.shields.io/pypi/pyversions/pystrukts.svg)](https://pypi.org/project/pystrukts)
77

88

9-
A Python library for data structures. This library provides more than 20 advanced data structures not available in the Python standard library.
9+
A Python library for data structures. This library provides more than 10 advanced data structures not available in the Python standard library.
1010

1111
## Installation
1212

@@ -21,23 +21,12 @@ pip install pystrukts
2121
- [x] Doubly Linked Node
2222
- [x] Doubly Linked List
2323
- [x] Circular Linked List (Single, Double)
24-
- [ ] Skip List
2524

2625
### Trees
2726
- [x] Binary Tree
2827
- [x] Generic Tree
2928
- [x] Binary Search Tree
30-
- [ ] Ternary Search Tree
31-
- [ ] Suffix Tree
32-
- [ ] AVL Tree
33-
- [ ] Red-Black Tree
34-
- [ ] Splay Tree
35-
- [ ] B-Tree
36-
- [ ] B+ Tree
37-
- [ ] R-Tree
38-
- [ ] Interval Tree
39-
- [ ] Segment Tree
40-
- [ ] Fenwick Tree
29+
- [x] Fenwick Trees (RUPQ, RURQ)
4130
- [x] Trie
4231

4332
### Graphs

pystrukts/tree/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,5 @@
33
from .trie import Trie, TrieNode
44
from .generic_tree import GenericTree
55
from .binary_tree import BinaryTree
6-
from .binary_search_tree import BinarySearchTree
6+
from .binary_search_tree import BinarySearchTree
7+
from .fenwick_tree import FenwickTree, RUPQ, RURQ

pystrukts/tree/avl.py

Whitespace-only changes.

pystrukts/tree/fenwick_tree.py

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
'''
2+
Fenwick Tree implementation in Python.
3+
4+
This module contains the implementation of Fenwick Trees, as well as two extensions of it:
5+
- Range Update Point Query (RUPQ)
6+
- Range Update Range Query (RURQ)
7+
'''
8+
9+
from typing import List
10+
11+
class FenwickTree:
12+
'''
13+
Fenwick Tree implementation in Python.
14+
15+
Attributes:
16+
- n: int - the size of the Fenwick Tree
17+
- ftree: list - the Fenwick Tree itself
18+
19+
Methods:
20+
- lsone(s: int) -> int: returns the least significant bit of s
21+
- query(i: int, j: int) -> int: returns the sum of the elements in the range [i, j]
22+
- update(i: int, v: int): updates the element at index i with value v
23+
- select(k: int) -> int: returns the index of the k-th element in the Fenwick Tree
24+
'''
25+
def __init__(self, f: List[int]):
26+
self.n = len(f)
27+
self.ftree = [0] * (self.n + 1)
28+
29+
for i in range(1, self.n + 1):
30+
self.ftree[i] += f[i - 1]
31+
if i + self._lsone(i) <= self.n:
32+
self.ftree[i + self._lsone(i)] += self.ftree[i]
33+
34+
def _lsone(self, s: int) -> int:
35+
'''
36+
Returns the least significant bit of s.
37+
38+
Args:
39+
s (int): The number to get the least significant bit from.
40+
41+
Returns:
42+
out (int): The least significant bit of s.
43+
'''
44+
return s & -s
45+
46+
def query(self, i: int, j: int) -> int:
47+
'''
48+
Queries the Fenwick Tree for the sum of the elements in the range [i, j].
49+
50+
Args:
51+
i (int): The lower bound of the range.
52+
j (int): The upper bound of the range.
53+
54+
Returns:
55+
s (int): The sum of the elements in the range [i, j].
56+
'''
57+
if i > 1:
58+
return self.query(1, j) - self.query(1, i - 1)
59+
60+
s = 0
61+
while j > 0:
62+
s += self.ftree[j]
63+
j -= self._lsone(j)
64+
65+
return s
66+
67+
def update(self, i: int, v: int):
68+
'''
69+
Updates the element at index i with value v.
70+
71+
Args:
72+
i (int): The index of the element to update.
73+
v (int): The new value of the element.
74+
'''
75+
76+
while i <= self.n:
77+
self.ftree[i] += v
78+
i += self._lsone(i)
79+
80+
def select(self, k: int) -> int:
81+
'''
82+
Returns the index of the k-th element in the Fenwick Tree.
83+
84+
Args:
85+
k (int): The index of the element to return.
86+
87+
Returns:
88+
i (int): The index of the k-th element in the Fenwick Tree.
89+
'''
90+
91+
p = 1
92+
while p*2 <= self.n:
93+
p *= 2
94+
95+
i = 0
96+
while p > 0:
97+
if k > self.ftree[i + p]:
98+
k -= self.ftree[i + p]
99+
i += p
100+
p //= 2
101+
102+
return i + 1
103+
104+
class RUPQ:
105+
'''
106+
Range Update Point Query (RUPQ) implementation.
107+
108+
This class extends the Fenwick Tree to support range updates and point queries.
109+
110+
Attributes:
111+
- ftree: FenwickTree - the Fenwick Tree used for the range updates
112+
113+
Methods:
114+
- query(i: int) -> int: returns the value of the element at index i
115+
- update(i: int, j: int, v: int): updates the elements in the range [i, j] with value v
116+
'''
117+
def __init__(self, n: int):
118+
self.ftree = FenwickTree([0] * n)
119+
120+
def query(self, i: int) -> int:
121+
'''
122+
Queries the Fenwick Tree for the value of the element at index i.
123+
124+
Args:
125+
i (int): The index of the element to query.
126+
127+
Returns:
128+
s (int): The value of the element at index i.
129+
'''
130+
131+
return self.ftree.query(1, i)
132+
133+
def update(self, i: int, j: int, v: int):
134+
'''
135+
Updates the elements in the range [i, j] with value v.
136+
137+
Args:
138+
i (int): The lower bound of the range.
139+
j (int): The upper bound of the range.
140+
v (int): The value to update the elements with.
141+
'''
142+
143+
self.ftree.update(i, v)
144+
self.ftree.update(j + 1, -v)
145+
146+
class RURQ:
147+
'''
148+
Range Update Range Query (RURQ) implementation.
149+
150+
This class extends the Fenwick Tree to support range updates and range queries.
151+
152+
Attributes:
153+
- f: FenwickTree - the Fenwick Tree used for the range updates
154+
- r: RUPQ - the RUPQ used for the range queries
155+
156+
Methods:
157+
- query(i: int, j: int) -> int: returns the sum of the elements in the range [i, j]
158+
- update(i: int, j: int, v: int): updates the elements in the range [i, j] with value v
159+
'''
160+
161+
def __init__(self, n: int):
162+
self.ftree = FenwickTree([0] * n)
163+
self.rupq = RUPQ(n)
164+
165+
def query(self, i: int, j: int) -> int:
166+
'''
167+
Queries the Fenwick Tree for the sum of the elements in the range [i, j].
168+
169+
Args:
170+
i (int): The lower bound of the range.
171+
j (int): The upper bound of the range.
172+
173+
Returns:
174+
s (int): The sum of the elements in the range [i, j].
175+
'''
176+
177+
if i > 1:
178+
return self.query(1, j) - self.query(1, i - 1)
179+
return self.rupq.query(j) * j - self.ftree.query(1, j)
180+
181+
def update(self, i: int, j: int, v: int):
182+
'''
183+
Updates the elements in the range [i, j] with value v.
184+
185+
Args:
186+
i (int): The lower bound of the range.
187+
j (int): The upper bound of the range.
188+
v (int): The value to update the elements with.
189+
'''
190+
191+
self.rupq.update(i, j, v)
192+
self.ftree.update(i, v * (i - 1))
193+
self.ftree.update(j + 1, -1 * v * j)

tests/test_fenwick_tree.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import pytest
2+
from pystrukts.tree import FenwickTree, RUPQ, RURQ
3+
4+
5+
def test_initialization_ftree():
6+
f = [0, 1, 0, 1, 2, 3, 2, 1, 1, 0]
7+
ft = FenwickTree(f)
8+
assert ft.n == 10
9+
assert ft.ftree == [0, 0, 1, 0, 2, 2, 5, 2, 10, 1, 1]
10+
assert ft.query(1, 6) == 7
11+
assert ft.query(1, 3) == 1
12+
assert ft.select(7) == 6
13+
ft.update(5, 1)
14+
assert ft.query(1, 10) == 12
15+
16+
def test_rupq():
17+
r = RUPQ(10)
18+
r.update(2, 9, 7)
19+
r.update(6, 7, 3)
20+
assert r.query(1) == 0
21+
assert r.query(2) == 7
22+
assert r.query(3) == 7
23+
assert r.query(4) == 7
24+
assert r.query(5) == 7
25+
assert r.query(6) == 10
26+
assert r.query(7) == 10
27+
assert r.query(8) == 7
28+
assert r.query(9) == 7
29+
assert r.query(10) == 0
30+
31+
def test_rurq():
32+
r = RURQ(10)
33+
r.update(2, 9, 7)
34+
r.update(6, 7, 3)
35+
assert r.query(3, 5) == 21
36+
assert r.query(7, 8) == 17

0 commit comments

Comments
 (0)