diff --git a/sqrt/sqrtdecomposition.go b/sqrt/sqrtdecomposition.go new file mode 100644 index 000000000..34e26fc4a --- /dev/null +++ b/sqrt/sqrtdecomposition.go @@ -0,0 +1,102 @@ +// Package sqrt contains algorithms and data structures that contains a √n in their complexity +package sqrt + +import "math" + +// Sqrt (or Square Root) Decomposition is a technique used for query an array and perform updates +// Inside this package is described its most simple data structure, you can find more at: https://cp-algorithms.com/data_structures/sqrt_decomposition.html +// +// Formally, You can use SqrtDecomposition only if: +// +// Given a function $Query:E_1,...,E_n\rightarrow Q$ +// +// if $\exist unionQ:Q,Q\rightarrow Q$ +// +// s.t. +// +// - $\forall n\in \N > 1, 1\le i 0, E_1,..., E_n\in E \\ query(E_1,...,E_{new},..., E_n)=updateQ(query(E_1,...,E_{old},...,E_n), indexof(E_{old}), E_{new})$ +type SqrtDecomposition[E any, Q any] struct { + querySingleElement func(element E) Q + unionQ func(q1 Q, q2 Q) Q + updateQ func(oldQ Q, oldE E, newE E) (newQ Q) + + elements []E + blocks []Q + blockSize uint64 +} + +// Create a new SqrtDecomposition instance with the parameters as specified by SqrtDecomposition comment +// Assumptions: +// - len(elements) > 0 +func NewSqrtDecomposition[E any, Q any]( + elements []E, + querySingleElement func(element E) Q, + unionQ func(q1 Q, q2 Q) Q, + updateQ func(oldQ Q, oldE E, newE E) (newQ Q), +) *SqrtDecomposition[E, Q] { + sqrtDec := &SqrtDecomposition[E, Q]{ + querySingleElement: querySingleElement, + unionQ: unionQ, + updateQ: updateQ, + elements: elements, + } + sqrt := math.Sqrt(float64(len(sqrtDec.elements))) + blockSize := uint64(sqrt) + numBlocks := uint64(math.Ceil(float64(len(elements)) / float64(blockSize))) + sqrtDec.blocks = make([]Q, numBlocks) + for i := uint64(0); i < uint64(len(elements)); i++ { + if i%blockSize == 0 { + sqrtDec.blocks[i/blockSize] = sqrtDec.querySingleElement(elements[i]) + } else { + sqrtDec.blocks[i/blockSize] = sqrtDec.unionQ(sqrtDec.blocks[i/blockSize], sqrtDec.querySingleElement(elements[i])) + } + } + sqrtDec.blockSize = blockSize + return sqrtDec +} + +// Performs a query from index start to index end (non included) +// Assumptions: +// - start < end +// - start and end are valid +func (s *SqrtDecomposition[E, Q]) Query(start uint64, end uint64) Q { + firstIndexNextBlock := ((start / s.blockSize) + 1) * s.blockSize + q := s.querySingleElement(s.elements[start]) + if firstIndexNextBlock > end { // if in same block + start++ + for start < end { + q = s.unionQ(q, s.querySingleElement(s.elements[start])) + start++ + } + } else { + // left side + start++ + for start < firstIndexNextBlock { + q = s.unionQ(q, s.querySingleElement(s.elements[start])) + start++ + } + + //middle part + endBlock := end / s.blockSize + for i := firstIndexNextBlock / s.blockSize; i < endBlock; i++ { + q = s.unionQ(q, s.blocks[i]) + } + + // right part + for i := endBlock * s.blockSize; i < end; i++ { + q = s.unionQ(q, s.querySingleElement(s.elements[i])) + } + } + return q +} + +// Assumptions: +// - index is valid +func (s *SqrtDecomposition[E, Q]) Update(index uint64, newElement E) { + i := index / s.blockSize + s.blocks[i] = s.updateQ(s.blocks[i], s.elements[index], newElement) + s.elements[index] = newElement +} diff --git a/sqrt/sqrtdecomposition_test.go b/sqrt/sqrtdecomposition_test.go new file mode 100644 index 000000000..7a293c802 --- /dev/null +++ b/sqrt/sqrtdecomposition_test.go @@ -0,0 +1,80 @@ +package sqrt_test + +import ( + "github.com/TheAlgorithms/Go/sqrt" + "testing" +) + +// Query interval +type query struct { + firstIndex uint64 + lastIndex uint64 +} + +type update struct { + index uint64 + value int +} + +func TestSqrtDecomposition(t *testing.T) { + var sqrtDecompositionTestData = []struct { + description string + array []int + updates []update + queries []query + expected []int + }{ + { + description: "test 1-sized array", + array: []int{1}, + queries: []query{{0, 1}}, + expected: []int{1}, + }, + { + description: "test array with size 5", + array: []int{1, 2, 3, 4, 5}, + queries: []query{{0, 5}, {0, 2}, {2, 4}}, + expected: []int{15, 3, 7}, + }, + { + description: "test array with size 5 and updates", + array: []int{1, 2, 3, 4, 5}, + updates: []update{{index: 1, value: 3}, + {index: 2, value: 4}}, + queries: []query{{0, 5}, {0, 2}, {2, 4}}, + expected: []int{17, 4, 8}, + }, + { + description: "test array with size 11 and updates", + array: []int{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1}, + updates: []update{{index: 2, value: 2}, + {index: 3, value: 3}, + {index: 6, value: 6}}, + queries: []query{{3, 5}, {7, 8}, {3, 7}, {0, 10}}, + expected: []int{4, 1, 11, 18}, + }, + } + for _, test := range sqrtDecompositionTestData { + t.Run(test.description, func(t *testing.T) { + s := sqrt.NewSqrtDecomposition(test.array, + func(e int) int { return e }, + func(q1, q2 int) int { return q1 + q2 }, + func(q, a, b int) int { return q - a + b }, + ) + + for i := 0; i < len(test.updates); i++ { + s.Update(test.updates[i].index, test.updates[i].value) + } + + for i := 0; i < len(test.queries); i++ { + result := s.Query(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) + } + } + + }) + } +}