diff --git a/structure/fenwicktree/fenwicktree.go b/structure/fenwicktree/fenwicktree.go new file mode 100644 index 000000000..915d7b085 --- /dev/null +++ b/structure/fenwicktree/fenwicktree.go @@ -0,0 +1,60 @@ +// Fenwick Tree Data Structure for efficient range queries on an array of integers. +// Also known as Binary Indexed Tree. It can query the sum of any range of the array and +// can update the array at a specific position by adding a value to it (point update). +// Build: O(N) +// Query: O(log(N)) +// Update: O(log(N)) +// reference: https://brilliant.org/wiki/fenwick-tree/ +package fenwicktree + +// FenwickTree represents the data structure of the Fenwick Tree +type FenwickTree struct { + n int // n: Size of the input array. + array []int // array: the input array on which queries are made. + bit []int // bit: store the sum of ranges. +} + +// NewFenwickTree creates a new Fenwick tree, initializes bit with +// the values of the array. Note that the queries and updates should have +// one based indexing. +func NewFenwickTree(array []int) *FenwickTree { + newArray := []int{0} // Appending a 0 to the beginning as this implementation uses 1 based indexing + fenwickTree := &FenwickTree{ + n: len(array), + array: append(newArray, array...), + bit: append(newArray, array...), + } + for i := 1; i < fenwickTree.n; i++ { + nextPos := i + (i & -i) + if nextPos <= fenwickTree.n { + fenwickTree.bit[nextPos] += fenwickTree.bit[i] + } + } + return fenwickTree +} + +// PrefixSum returns the sum of the prefix ending at position pos. +func (f *FenwickTree) PrefixSum(pos int) int { + if pos > f.n { + return 0 + } + prefixSum := 0 + for i := pos; i > 0; i -= (i & -i) { + prefixSum += f.bit[i] + } + return prefixSum +} + +// RangeSum returns the sum of the elements in the range l to r +// both inclusive. +func (f *FenwickTree) RangeSum(l int, r int) int { + return f.PrefixSum(r) - f.PrefixSum(l-1) +} + +// Add Adds value to the element at position pos of the array +// and recomputes the range sums. +func (f *FenwickTree) Add(pos int, value int) { + for i := pos; i <= f.n; i += (i & -i) { + f.bit[i] += value + } +} diff --git a/structure/fenwicktree/fenwicktree_test.go b/structure/fenwicktree/fenwicktree_test.go new file mode 100644 index 000000000..99c9f12fc --- /dev/null +++ b/structure/fenwicktree/fenwicktree_test.go @@ -0,0 +1,81 @@ +package fenwicktree_test + +import ( + "github.com/TheAlgorithms/Go/structure/fenwicktree" + "testing" +) + +type query struct { + queryType string + firstIndex int // firstIndex and lastIndex are same for point queries + lastIndex int +} + +type update struct { + pos int + value int +} + +func TestFenwickTree(t *testing.T) { + var fenwickTreeTestData = []struct { + description string + array []int + updates []update + queries []query + expected []int + }{ + { + description: "test empty array", + array: []int{}, + queries: []query{{"point", 1, 1}}, + expected: []int{0}, + }, + { + description: "test array with size 5", + array: []int{1, 2, 3, 4, 5}, + queries: []query{{"range", 1, 5}, {"range", 1, 3}, {"range", 3, 5}}, + expected: []int{15, 6, 12}, + }, + { + description: "test array with size 5, single index updates and range queries", + array: []int{1, 2, 3, 4, 5}, + updates: []update{{pos: 2, value: 2}, {pos: 3, value: 3}}, + queries: []query{{"range", 1, 5}, {"range", 1, 3}, {"range", 3, 5}}, + expected: []int{20, 11, 15}, + }, + { + description: "test array with size 5, single index updates and point queries", + array: []int{1, 2, 3, 4, 5}, + updates: []update{{pos: 2, value: 2}, {pos: 3, value: 3}}, + queries: []query{{"point", 3, 3}, {"point", 1, 1}, {"point", 5, 5}}, + expected: []int{11, 1, 20}, + }, + } + + for _, test := range fenwickTreeTestData { + t.Run(test.description, func(t *testing.T) { + fenwickTree := fenwicktree.NewFenwickTree(test.array) + + for i := 0; i < len(test.updates); i++ { + fenwickTree.Add(test.updates[i].pos, test.updates[i].value) + } + + for i := 0; i < len(test.queries); i++ { + + var result int + + if test.queries[i].queryType == "point" { + result = fenwickTree.PrefixSum(test.queries[i].firstIndex) + } else { + result = fenwickTree.RangeSum(test.queries[i].firstIndex, test.queries[i].lastIndex) + } + + if result != test.expected[i] { + t.Logf("FAIL: %s", test.description) + t.Fatalf("Expected result: %d\nFound: %d\n", test.expected[i], result) + } + } + }) + } + +}