diff --git a/src/data-structures/tree/fenwick-tree/FenwickTree.js b/src/data-structures/tree/fenwick-tree/FenwickTree.js new file mode 100644 index 00000000..bcce8727 --- /dev/null +++ b/src/data-structures/tree/fenwick-tree/FenwickTree.js @@ -0,0 +1,49 @@ +export default class FenwickTree { + /** + * Constructor creates empty fenwick tree of size 'size', + * however, array size is size+1, because index is 1-based + * @param {number} [size] + */ + constructor(size) { + this.n = size; + this.arr = []; + for (let i = 0; i <= size; i += 1) this.arr.push(0); + } + + /** + * Adds v to index x + * @param {number} [x] + * @param {number} [v] + */ + update(x, v) { + if (x < 1 || x > this.n) return; + for (let i = x; i <= this.n; i += (i & -i)) { + this.arr[i] += v; + } + } + + /** + * query sum from index 1 to x + * @param {number} [x] + * @return {number} sum + */ + query(x) { + if (x > this.n) return this.query(this.n); + let ret = 0; + for (let i = x; i > 0; i -= (i & -i)) { + ret += this.arr[i]; + } + return ret; + } + + /** + * query sum from index l to r + * @param {number} [l] + * @param {number} [r] + * @return {number} + */ + queryRange(l, r) { + if (l > r) return 0; + return this.query(r) - this.query(l - 1); + } +} diff --git a/src/data-structures/tree/fenwick-tree/README.md b/src/data-structures/tree/fenwick-tree/README.md new file mode 100644 index 00000000..2b7c9905 --- /dev/null +++ b/src/data-structures/tree/fenwick-tree/README.md @@ -0,0 +1,8 @@ +# Binary Indexed Tree / Fenwick Tree + +A simple data structure that supports fast range queries in an array. However, it is usually only valid for reversible operations, like addition and subtraction + +This implementation uses the basic range sum query and point update. All the indexes are 1-based + +- [Wikipedia](https://en.wikipedia.org/wiki/Fenwick_tree) +- [Geeksforgeeks](https://www.geeksforgeeks.org/binary-indexed-tree-or-fenwick-tree-2/) \ No newline at end of file diff --git a/src/data-structures/tree/fenwick-tree/__test__/FenwickTree.test.js b/src/data-structures/tree/fenwick-tree/__test__/FenwickTree.test.js new file mode 100644 index 00000000..4099cf30 --- /dev/null +++ b/src/data-structures/tree/fenwick-tree/__test__/FenwickTree.test.js @@ -0,0 +1,37 @@ +import FenwickTree from '../FenwickTree'; + +describe('FenwickTree', () => { + it('should create empty fenwick tree of correct size', () => { + const tree1 = new FenwickTree(5); + expect(tree1.arr.length).toBe(5 + 1); + + for (let i = 0; i < 5; i += 1) { + expect(tree1.arr[i]).toBe(0); + } + + const tree2 = new FenwickTree(50); + expect(tree2.arr.length).toBe(50 + 1); + }); + + it('should correctly execute queries', () => { + const tree = new FenwickTree(5); + + tree.update(1, 4); + tree.update(3, 7); + + expect(tree.query(1)).toBe(4); + expect(tree.query(3)).toBe(11); + expect(tree.query(5)).toBe(11); + expect(tree.queryRange(2, 3)).toBe(7); + + tree.update(2, 5); + expect(tree.query(5)).toBe(16); + + tree.update(1, 3); + expect(tree.queryRange(1, 1)).toBe(7); + expect(tree.query(5)).toBe(19); + expect(tree.queryRange(1, 5)).toBe(19); + + expect(tree.queryRange(5, 1)).toBe(0); // invalid test + }); +});