#### Sqrtdecomposition

// 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<n, E_1,..., E_n\in E \\ query(E_1,..., E_n)=unionQ(query(E_1,..., E_i), query(E_{i+1},...,E_n))$
//
// - (Only if you want use $update$ function)
// $\forall n\in \N > 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
}